################################################## # Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021 # ################################################## from __future__ import division from __future__ import print_function import os, math, random from collections import OrderedDict import numpy as np import pandas as pd import copy from functools import partial from typing import Optional, Text from qlib.utils import ( unpack_archive_with_buffer, save_multiple_parts_file, get_or_create_path, drop_nan_by_y_index, ) from qlib.log import get_module_logger import torch import torch.nn.functional as F import torch.optim as optim import torch.utils.data as th_data from log_utils import AverageMeter from utils import count_parameters from trade_models.transformers import DEFAULT_NET_CONFIG from trade_models.transformers import get_transformer from qlib.model.base import Model from qlib.data.dataset import DatasetH from qlib.data.dataset.handler import DataHandlerLP DEFAULT_OPT_CONFIG = dict( epochs=200, lr=0.001, batch_size=2000, early_stop=20, loss="mse", optimizer="adam", num_workers=4 ) class QuantTransformer(Model): """Transformer-based Quant Model""" def __init__(self, net_config=None, opt_config=None, metric="", GPU=0, seed=None, **kwargs): # Set logger. self.logger = get_module_logger("QuantTransformer") self.logger.info("QuantTransformer PyTorch version...") # set hyper-parameters. self.net_config = net_config or DEFAULT_NET_CONFIG self.opt_config = opt_config or DEFAULT_OPT_CONFIG self.metric = metric self.device = torch.device("cuda:{:}".format(GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu") self.seed = seed self.logger.info( "Transformer parameters setting:" "\nnet_config : {:}" "\nopt_config : {:}" "\nmetric : {:}" "\ndevice : {:}" "\nseed : {:}".format( self.net_config, self.opt_config, self.metric, self.device, self.seed, ) ) if self.seed is not None: random.seed(self.seed) np.random.seed(self.seed) torch.manual_seed(self.seed) if self.use_gpu: torch.cuda.manual_seed(self.seed) torch.cuda.manual_seed_all(self.seed) self.model = get_transformer(self.net_config) self.logger.info("model: {:}".format(self.model)) self.logger.info("model size: {:.3f} MB".format(count_parameters(self.model))) if self.opt_config["optimizer"] == "adam": self.train_optimizer = optim.Adam(self.model.parameters(), lr=self.opt_config["lr"]) elif self.opt_config["optimizer"] == "adam": self.train_optimizer = optim.SGD(self.model.parameters(), lr=self.opt_config["lr"]) else: raise NotImplementedError("optimizer {:} is not supported!".format(optimizer)) self.fitted = False self.model.to(self.device) @property def use_gpu(self): return self.device != torch.device("cpu") def loss_fn(self, pred, label): mask = ~torch.isnan(label) if self.opt_config["loss"] == "mse": return F.mse_loss(pred[mask], label[mask]) else: raise ValueError("unknown loss `{:}`".format(self.loss)) def metric_fn(self, pred, label): # the metric score : higher is better if self.metric == "" or self.metric == "loss": return -self.loss_fn(pred, label) else: raise ValueError("unknown metric `{:}`".format(self.metric)) def train_or_test_epoch(self, xloader, model, loss_fn, metric_fn, is_train, optimizer=None): if is_train: model.train() else: model.eval() score_meter, loss_meter = AverageMeter(), AverageMeter() for ibatch, (feats, labels) in enumerate(xloader): feats = feats.to(self.device, non_blocking=True) labels = labels.to(self.device, non_blocking=True) # forward the network preds = model(feats) loss = loss_fn(preds, labels) with torch.no_grad(): score = self.metric_fn(preds, labels) loss_meter.update(loss.item(), feats.size(0)) score_meter.update(score.item(), feats.size(0)) # optimize the network if is_train and optimizer is not None: optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_value_(model.parameters(), 3.0) optimizer.step() return loss_meter.avg, score_meter.avg def fit( self, dataset: DatasetH, save_path: Optional[Text] = None, ): def _prepare_dataset(df_data): return th_data.TensorDataset( torch.from_numpy(df_data["feature"].values).float(), torch.from_numpy(df_data["label"].values).squeeze().float(), ) def _prepare_loader(dataset, shuffle): return th_data.DataLoader( dataset, batch_size=self.opt_config["batch_size"], drop_last=False, pin_memory=True, num_workers=self.opt_config["num_workers"], shuffle=shuffle, ) df_train, df_valid, df_test = dataset.prepare( ["train", "valid", "test"], col_set=["feature", "label"], data_key=DataHandlerLP.DK_L, ) train_dataset, valid_dataset, test_dataset = ( _prepare_dataset(df_train), _prepare_dataset(df_valid), _prepare_dataset(df_test), ) train_loader, valid_loader, test_loader = ( _prepare_loader(train_dataset, True), _prepare_loader(valid_dataset, False), _prepare_loader(test_dataset, False), ) save_path = get_or_create_path(save_path, return_dir=True) self.logger.info("Fit procedure for [{:}] with save path={:}".format(self.__class__.__name__, save_path)) def _internal_test(ckp_epoch=None, results_dict=None): with torch.no_grad(): train_loss, train_score = self.train_or_test_epoch( train_loader, self.model, self.loss_fn, self.metric_fn, False, None ) valid_loss, valid_score = self.train_or_test_epoch( valid_loader, self.model, self.loss_fn, self.metric_fn, False, None ) test_loss, test_score = self.train_or_test_epoch( test_loader, self.model, self.loss_fn, self.metric_fn, False, None ) xstr = "train-score={:.6f}, valid-score={:.6f}, test-score={:.6f}".format( train_score, valid_score, test_score ) if ckp_epoch is not None and isinstance(results_dict, dict): results_dict["train"][ckp_epoch] = train_score results_dict["valid"][ckp_epoch] = valid_score results_dict["test"][ckp_epoch] = test_score return dict(train=train_score, valid=valid_score, test=test_score), xstr # Pre-fetch the potential checkpoints ckp_path = os.path.join(save_path, "{:}.pth".format(self.__class__.__name__)) if os.path.exists(ckp_path): ckp_data = torch.load(ckp_path) import pdb pdb.set_trace() else: stop_steps, best_score, best_epoch = 0, -np.inf, -1 start_epoch = 0 results_dict = dict(train=OrderedDict(), valid=OrderedDict(), test=OrderedDict()) _, eval_str = _internal_test(-1, results_dict) self.logger.info("Training from scratch, metrics@start: {:}".format(eval_str)) for iepoch in range(start_epoch, self.opt_config["epochs"]): self.logger.info( "Epoch={:03d}/{:03d} ::==>> Best valid @{:03d} ({:.6f})".format( iepoch, self.opt_config["epochs"], best_epoch, best_score ) ) train_loss, train_score = self.train_or_test_epoch( train_loader, self.model, self.loss_fn, self.metric_fn, True, self.train_optimizer ) self.logger.info("Training :: loss={:.6f}, score={:.6f}".format(train_loss, train_score)) current_eval_scores, eval_str = _internal_test(iepoch, results_dict) self.logger.info("Evaluating :: {:}".format(eval_str)) if current_eval_scores["valid"] > best_score: stop_steps, best_epoch, best_score = 0, iepoch, current_eval_scores["valid"] best_param = copy.deepcopy(self.model.state_dict()) else: stop_steps += 1 if stop_steps >= self.opt_config["early_stop"]: self.logger.info("early stop at {:}-th epoch, where the best is @{:}".format(iepoch, best_epoch)) break save_info = dict( net_config=self.net_config, opt_config=self.opt_config, net_state_dict=self.model.state_dict(), opt_state_dict=self.train_optimizer.state_dict(), best_param=best_param, stop_steps=stop_steps, best_score=best_score, best_epoch=best_epoch, start_epoch=iepoch + 1, ) torch.save(save_info, ckp_path) self.logger.info("The best score: {:.6f} @ {:02d}-th epoch".format(best_score, best_epoch)) self.model.load_state_dict(best_param) if self.use_gpu: torch.cuda.empty_cache() self.fitted = True def predict(self, dataset): if not self.fitted: raise ValueError("The model is not fitted yet!") x_test = dataset.prepare("test", col_set="feature") index = x_test.index self.model.eval() x_values = x_test.values sample_num, batch_size = x_values.shape[0], self.opt_config["batch_size"] preds = [] for begin in range(sample_num)[::batch_size]: if sample_num - begin < batch_size: end = sample_num else: end = begin + batch_size x_batch = torch.from_numpy(x_values[begin:end]).float().to(self.device) with torch.no_grad(): pred = self.model(x_batch).detach().cpu().numpy() preds.append(pred) return pd.Series(np.concatenate(preds), index=index)