auto-trading/notebooks/model_evaluation.py

36 lines
1.3 KiB
Python

# %% Import required packages
import torch
from src.models.transformer_model import TransformerModel
from src.models.rl_model import RLModel
from src.models.trading_agent import TradingAgent
from src.evaluation.evaluate import evaluate_trading_agent
from src.data.data_preprocessing import load_processed_data
# %% Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# %% Load processed data
data = load_processed_data('./data/processed/processed_data.csv')
# %% Initialize models
transformer_model = TransformerModel().to(device)
rl_model = RLModel().to(device)
trading_agent = TradingAgent(transformer_model, rl_model)
# %% Load model weights
transformer_model.load_state_dict(torch.load('./models/transformer_model.pth'))
rl_model.load_state_dict(torch.load('./models/rl_model.pth'))
# %% Evaluate the trading agent
trading_agent_results = evaluate_trading_agent(trading_agent, data)
# %% Display evaluation results
print("Total Profit: ", trading_agent_results['total_profit'])
print("Total Trades Made: ", trading_agent_results['total_trades'])
print("Successful Trades: ", trading_agent_results['successful_trades'])
# %% Save evaluation results
with open('./logs/evaluation_results.txt', 'w') as f:
for key, value in trading_agent_results.items():
f.write(f'{key}: {value}\n')