47 lines
1.5 KiB
Python
47 lines
1.5 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.training.train import train_transformer, train_rl
|
|
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)
|
|
|
|
# %% Train Transformer Model
|
|
# Set the appropriate hyperparameters
|
|
transformer_model_hyperparams = {
|
|
"epochs": 10,
|
|
"batch_size": 32,
|
|
"learning_rate": 0.001,
|
|
}
|
|
train_transformer(transformer_model, data, transformer_model_hyperparams)
|
|
|
|
# %% Save Transformer Model
|
|
torch.save(transformer_model.state_dict(), './models/transformer_model.pth')
|
|
|
|
# %% Train RL Model
|
|
# Set the appropriate hyperparameters
|
|
rl_model_hyperparams = {
|
|
"epochs": 500,
|
|
"batch_size": 32,
|
|
"learning_rate": 0.001,
|
|
"gamma": 0.99, # discount factor
|
|
"epsilon_start": 1.0, # exploration rate at the beginning
|
|
"epsilon_end": 0.01, # minimum exploration rate
|
|
"epsilon_decay": 0.995, # exponential decay rate for exploration probability
|
|
}
|
|
train_rl(trading_agent, data, rl_model_hyperparams)
|
|
|
|
# %% Save RL Model
|
|
torch.save(rl_model.state_dict(), './models/rl_model.pth')
|