liuyu/qlib
1
0
Fork 0
mirror of https://github.com/microsoft/qlib.git synced 2026-07-13 15:56:57 +08:00
Code Activity

Fix alstm model.

Browse Source
This commit is contained in:
lwwang1995
2020-11-27 11:03:44 +08:00
parent 55acac9fd5
commit c8355f9f18
1 changed files with 23 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

67
qlib/contrib/model/pytorch_alstm.py
View File

@@ -9,8 +9,10 @@ import os
import numpy as np import numpy as np
import pandas as pd import pandas as pd
import copy import copy
from ...utils import create_save_path from sklearn.metrics import roc_auc_score, mean_squared_error
from ...log import get_module_logger import logging
from ...utils import unpack_archive_with_buffer, save_multiple_parts_file, create_save_path, drop_nan_by_y_index
from ...log import get_module_logger, TimeInspector
import torch import torch
import torch.nn as nn import torch.nn as nn
@@ -51,7 +53,6 @@ class ALSTM(Model):
optimizer="adam", optimizer="adam",
GPU="0", GPU="0",
seed=0, seed=0,
rnn_type="GRU",
**kwargs **kwargs
): ):
# Set logger. # Set logger.
@@ -73,7 +74,6 @@ class ALSTM(Model):
self.visible_GPU = GPU self.visible_GPU = GPU
self.use_gpu = torch.cuda.is_available() self.use_gpu = torch.cuda.is_available()
self.seed = seed self.seed = seed
self.rnn_type = rnn_type
self.logger.info( self.logger.info(
"ALSTM parameters setting:" "ALSTM parameters setting:"
@@ -90,8 +90,7 @@ class ALSTM(Model):
"\nloss_type : {}" "\nloss_type : {}"
"\nvisible_GPU : {}" "\nvisible_GPU : {}"
"\nuse_GPU : {}" "\nuse_GPU : {}"
"\nseed : {}" "\nseed : {}".format(
"\nrnn_type : {}".format(
d_feat, d_feat,
hidden_size, hidden_size,
num_layers, num_layers,
@@ -106,24 +105,22 @@ class ALSTM(Model):
GPU, GPU,
self.use_gpu, self.use_gpu,
seed, seed,
self.rnn_type,
) )
) )
self.alstm_model = ALSTMModel( self.ALSTM_model = ALSTMModel(
d_feat=self.d_feat, hidden_size=self.hidden_size, num_layers=self.num_layers, dropout=self.dropout d_feat=self.d_feat, hidden_size=self.hidden_size, num_layers=self.num_layers, dropout=self.dropout
) )
if optimizer.lower() == "adam": if optimizer.lower() == "adam":
self.train_optimizer = optim.Adam(self.alstm_model.parameters(), lr=self.lr) self.train_optimizer = optim.Adam(self.ALSTM_model.parameters(), lr=self.lr)
elif optimizer.lower() == "gd": elif optimizer.lower() == "gd":
self.train_optimizer = optim.SGD(self.alstm_model.parameters(), lr=self.lr) self.train_optimizer = optim.SGD(self.ALSTM_model.parameters(), lr=self.lr)
else: else:
raise NotImplementedError("optimizer {} is not supported!".format(optimizer)) raise NotImplementedError("optimizer {} is not supported!".format(optimizer))
self._fitted = False self._fitted = False
if self.use_gpu: if self.use_gpu:
self.alstm_model.cuda() self.ALSTM_model.cuda()
# set the visible GPU # set the visible GPU
if self.visible_GPU: if self.visible_GPU:
os.environ["CUDA_VISIBLE_DEVICES"] = self.visible_GPU os.environ["CUDA_VISIBLE_DEVICES"] = self.visible_GPU
@@ -141,6 +138,7 @@ class ALSTM(Model):
raise ValueError("unknown loss `%s`" % self.loss) raise ValueError("unknown loss `%s`" % self.loss)
def metric_fn(self, pred, label): def metric_fn(self, pred, label):
mask = torch.isfinite(label) mask = torch.isfinite(label)
if self.metric == "" or self.metric == "loss": # use loss if self.metric == "" or self.metric == "loss": # use loss
@@ -148,12 +146,13 @@ class ALSTM(Model):
raise ValueError("unknown metric `%s`" % self.metric) raise ValueError("unknown metric `%s`" % self.metric)
def train_epoch(self, x_train, y_train): def train_epoch(self, x_train, y_train):
x_train_values = x_train.values x_train_values = x_train.values
y_train_values = np.squeeze(y_train.values) y_train_values = np.squeeze(y_train.values)
self.alstm_model.train() self.ALSTM_model.train()
indices = np.arange(len(x_train_values)) indices = np.arange(len(x_train_values))
np.random.shuffle(indices) np.random.shuffle(indices)
@@ -170,21 +169,21 @@ class ALSTM(Model):
feature = feature.cuda() feature = feature.cuda()
label = label.cuda() label = label.cuda()
pred = self.alstm_model(feature) pred = self.ALSTM_model(feature)
loss = self.loss_fn(pred, label) loss = self.loss_fn(pred, label)
self.train_optimizer.zero_grad() self.train_optimizer.zero_grad()
loss.backward() loss.backward()
torch.nn.utils.clip_grad_value_(self.alstm_model.parameters(), 3.0) torch.nn.utils.clip_grad_value_(self.ALSTM_model.parameters(), 3.0)
self.train_optimizer.step() self.train_optimizer.step()
def test_epoch(self, data_x, data_y): def test_epoch(self, data_x, data_y):
# prepare testing data # prepare training data
x_values = data_x.values x_values = data_x.values
y_values = np.squeeze(data_y.values) y_values = np.squeeze(data_y.values)
self.alstm_model.eval() self.ALSTM_model.eval()
scores = [] scores = []
losses = [] losses = []
@@ -203,7 +202,7 @@ class ALSTM(Model):
feature = feature.cuda() feature = feature.cuda()
label = label.cuda() label = label.cuda()
pred = self.alstm_model(feature) pred = self.ALSTM_model(feature)
loss = self.loss_fn(pred, label) loss = self.loss_fn(pred, label)
losses.append(loss.item()) losses.append(loss.item())
@@ -230,6 +229,7 @@ class ALSTM(Model):
if save_path == None: if save_path == None:
save_path = create_save_path(save_path) save_path = create_save_path(save_path)
stop_steps = 0 stop_steps = 0
train_loss = 0
best_score = -np.inf best_score = -np.inf
best_epoch = 0 best_epoch = 0
evals_result["train"] = [] evals_result["train"] = []
@@ -254,7 +254,7 @@ class ALSTM(Model):
best_score = val_score best_score = val_score
stop_steps = 0 stop_steps = 0
best_epoch = step best_epoch = step
best_param = copy.deepcopy(self.alstm_model.state_dict()) best_param = copy.deepcopy(self.ALSTM_model.state_dict())
else: else:
stop_steps += 1 stop_steps += 1
if stop_steps >= self.early_stop: if stop_steps >= self.early_stop:
@@ -262,7 +262,7 @@ class ALSTM(Model):
break break
self.logger.info("best score: %.6lf @ %d" % (best_score, best_epoch)) self.logger.info("best score: %.6lf @ %d" % (best_score, best_epoch))
self.alstm_model.load_state_dict(best_param) self.ALSTM_model.load_state_dict(best_param)
torch.save(best_param, save_path) torch.save(best_param, save_path)
if self.use_gpu: if self.use_gpu:
@@ -274,7 +274,7 @@ class ALSTM(Model):
x_test = dataset.prepare("test", col_set="feature") x_test = dataset.prepare("test", col_set="feature")
index = x_test.index index = x_test.index
self.alstm_model.eval() self.ALSTM_model.eval()
x_values = x_test.values x_values = x_test.values
sample_num = x_values.shape[0] sample_num = x_values.shape[0]
preds = [] preds = []
@@ -293,36 +293,15 @@ class ALSTM(Model):
with torch.no_grad(): with torch.no_grad():
if self.use_gpu: if self.use_gpu:
pred = self.alstm_model(x_batch).detach().cpu().numpy() pred = self.ALSTM_model(x_batch).detach().cpu().numpy()
else: else:
pred = self.alstm_model(x_batch).detach().numpy() pred = self.ALSTM_model(x_batch).detach().numpy()
preds.append(pred) preds.append(pred)
return pd.Series(np.concatenate(preds), index=index) return pd.Series(np.concatenate(preds), index=index)
class GRUModel(nn.Module):
def __init__(self, d_feat=6, hidden_size=64, num_layers=2, dropout=0.0):
super().__init__()
self.rnn = nn.GRU(
input_size=d_feat,
hidden_size=hidden_size,
num_layers=num_layers,
batch_first=True,
dropout=dropout,
)
self.fc_out = nn.Linear(hidden_size, 1)
self.d_feat = d_feat
def forward(self, x):
x = x.reshape(len(x), self.d_feat, -1) # [N, F, T]
x = x.permute(0, 2, 1) # [N, T, F]
out, _ = self.rnn(x)
return self.fc_out(out[:, -1, :]).squeeze()
class ALSTMModel(nn.Module): class ALSTMModel(nn.Module):
def __init__(self, d_feat=6, hidden_size=64, num_layers=2, dropout=0.0, rnn_type="GRU"): def __init__(self, d_feat=6, hidden_size=64, num_layers=2, dropout=0.0, rnn_type="GRU"):
super().__init__() super().__init__()