Merge pull request microsoft#378 from D-X-Y/main

Add MultiSegRecord and add segment kwargs in model.pred

Author: you-n-g

qlib/contrib/model/__init__.py

@@ -0,0 +1,39 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+try:
+    from .catboost_model import CatBoostModel
+except ModuleNotFoundError:
+    CatBoostModel = None
+    print("Please install necessary libs for CatBoostModel.")
+try:
+    from .double_ensemble import DEnsembleModel
+    from .gbdt import LGBModel
+except ModuleNotFoundError:
+    DEnsembleModel, LGBModel = None, None
+    print("Please install necessary libs for DEnsembleModel and LGBModel, such as lightgbm.")
+try:
+    from .xgboost import XGBModel
+except ModuleNotFoundError:
+    XGBModel = None
+    print("Please install necessary libs for XGBModel, such as xgboost.")
+try:
+    from .linear import LinearModel
+except ModuleNotFoundError:
+    LinearModel = None
+    print("Please install necessary libs for LinearModel, such as scipy and sklearn.")
+# import pytorch models
+try:
+    from .pytorch_alstm import ALSTM
+    from .pytorch_gats import GATs
+    from .pytorch_gru import GRU
+    from .pytorch_lstm import LSTM
+    from .pytorch_nn import DNNModelPytorch
+    from .pytorch_tabnet import TabnetModel
+    from .pytorch_sfm import SFM_Model
+
+    pytorch_classes = (ALSTM, GATs, GRU, LSTM, DNNModelPytorch, TabnetModel, SFM_Model)
+except ModuleNotFoundError:
+    pytorch_classes = ()
+    print("Please install necessary libs for PyTorch models.")
+
+all_model_classes = (CatBoostModel, DEnsembleModel, LGBModel, XGBModel, LinearModel) + pytorch_classes

qlib/contrib/model/catboost_model.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 import pandas as pd
+from typing import Text, Union
 from catboost import Pool, CatBoost
 from catboost.utils import get_gpu_device_count
 
@@ -62,10 +63,10 @@ def fit(
         evals_result["train"] = list(evals_result["learn"].values())[0]
         evals_result["valid"] = list(evals_result["validation"].values())[0]
 
-    def predict(self, dataset):
+    def predict(self, dataset: DatasetH, segment: Union[Text, slice] = "test"):
         if self.model is None:
             raise ValueError("model is not fitted yet!")
-        x_test = dataset.prepare("test", col_set="feature")
+        x_test = dataset.prepare(segment, col_set="feature", data_key=DataHandlerLP.DK_I)
         return pd.Series(self.model.predict(x_test.values), index=x_test.index)
 
 

qlib/contrib/model/double_ensemble.py

@@ -4,7 +4,7 @@
 import lightgbm as lgb
 import numpy as np
 import pandas as pd
-
+from typing import Text, Union
 from ...model.base import Model
 from ...data.dataset import DatasetH
 from ...data.dataset.handler import DataHandlerLP
@@ -40,6 +40,10 @@ def __init__(
         self.bins_sr = bins_sr
         self.bins_fs = bins_fs
         self.decay = decay
+        if sample_ratios is None:  # the default values for sample_ratios
+            sample_ratios = [0.8, 0.7, 0.6, 0.5, 0.4]
+        if sub_weights is None:  # the default values for sub_weights
+            sub_weights = [1.0, 0.2, 0.2, 0.2, 0.2, 0.2]
         if not len(sample_ratios) == bins_fs:
             raise ValueError("The length of sample_ratios should be equal to bins_fs.")
         self.sample_ratios = sample_ratios
@@ -228,10 +232,10 @@ def retrieve_loss_curve(self, model, df_train, features):
             raise ValueError("not implemented yet")
         return loss_curve
 
-    def predict(self, dataset):
+    def predict(self, dataset: DatasetH, segment: Union[Text, slice] = "test"):
         if self.ensemble is None:
             raise ValueError("model is not fitted yet!")
-        x_test = dataset.prepare("test", col_set="feature", data_key=DataHandlerLP.DK_I)
+        x_test = dataset.prepare(segment, col_set="feature", data_key=DataHandlerLP.DK_I)
         pred = pd.Series(np.zeros(x_test.shape[0]), index=x_test.index)
         for i_sub, submodel in enumerate(self.ensemble):
             feat_sub = self.sub_features[i_sub]

qlib/contrib/model/gbdt.py

@@ -4,7 +4,7 @@
 import numpy as np
 import pandas as pd
 import lightgbm as lgb
-
+from typing import Text, Union
 from ...model.base import ModelFT
 from ...data.dataset import DatasetH
 from ...data.dataset.handler import DataHandlerLP
@@ -61,10 +61,10 @@ def fit(
         evals_result["train"] = list(evals_result["train"].values())[0]
         evals_result["valid"] = list(evals_result["valid"].values())[0]
 
-    def predict(self, dataset):
+    def predict(self, dataset: DatasetH, segment: Union[Text, slice] = "test"):
         if self.model is None:
             raise ValueError("model is not fitted yet!")
-        x_test = dataset.prepare("test", col_set="feature", data_key=DataHandlerLP.DK_I)
+        x_test = dataset.prepare(segment, col_set="feature", data_key=DataHandlerLP.DK_I)
         return pd.Series(self.model.predict(x_test.values), index=x_test.index)
 
     def finetune(self, dataset: DatasetH, num_boost_round=10, verbose_eval=20):

qlib/contrib/model/linear.py

@@ -3,7 +3,7 @@
 
 import numpy as np
 import pandas as pd
-
+from typing import Text, Union
 from scipy.optimize import nnls
 from sklearn.linear_model import LinearRegression, Ridge, Lasso
 
@@ -84,8 +84,8 @@ def _fit_nnls(self, X, y):
             self.coef_ = coef
             self.intercept_ = 0.0
 
-    def predict(self, dataset):
+    def predict(self, dataset: DatasetH, segment: Union[Text, slice] = "test"):
         if self.coef_ is None:
             raise ValueError("model is not fitted yet!")
-        x_test = dataset.prepare("test", col_set="feature", data_key=DataHandlerLP.DK_I)
+        x_test = dataset.prepare(segment, col_set="feature", data_key=DataHandlerLP.DK_I)
         return pd.Series(x_test.values @ self.coef_ + self.intercept_, index=x_test.index)

qlib/contrib/model/pytorch_alstm.py

@@ -8,13 +8,9 @@
 import os
 import numpy as np
 import pandas as pd
+from typing import Text, Union
 import copy
-from ...utils import (
-    unpack_archive_with_buffer,
-    save_multiple_parts_file,
-    get_or_create_path,
-    drop_nan_by_y_index,
-)
+from ...utils import get_or_create_path
 from ...log import get_module_logger
 
 import torch
@@ -273,11 +269,11 @@ def fit(
         if self.use_gpu:
             torch.cuda.empty_cache()
 
-    def predict(self, dataset):
+    def predict(self, dataset: DatasetH, segment: Union[Text, slice] = "test"):
         if not self.fitted:
             raise ValueError("model is not fitted yet!")
 
-        x_test = dataset.prepare("test", col_set="feature")
+        x_test = dataset.prepare(segment, col_set="feature", data_key=DataHandlerLP.DK_I)
         index = x_test.index
         self.ALSTM_model.eval()
         x_values = x_test.values

qlib/contrib/model/pytorch_alstm_ts.py

@@ -8,13 +8,9 @@
 import os
 import numpy as np
 import pandas as pd
+from typing import Text, Union
 import copy
-from ...utils import (
-    unpack_archive_with_buffer,
-    save_multiple_parts_file,
-    get_or_create_path,
-    drop_nan_by_y_index,
-)
+from ...utils import get_or_create_path
 from ...log import get_module_logger
 
 import torch
@@ -264,11 +260,11 @@ def fit(
         if self.use_gpu:
             torch.cuda.empty_cache()
 
-    def predict(self, dataset):
+    def predict(self, dataset: DatasetH, segment: Union[Text, slice] = "test"):
         if not self.fitted:
             raise ValueError("model is not fitted yet!")
 
-        dl_test = dataset.prepare("test", col_set=["feature", "label"], data_key=DataHandlerLP.DK_I)
+        dl_test = dataset.prepare(segment, col_set=["feature", "label"], data_key=DataHandlerLP.DK_I)
         dl_test.config(fillna_type="ffill+bfill")
         test_loader = DataLoader(dl_test, batch_size=self.batch_size, num_workers=self.n_jobs)
         self.ALSTM_model.eval()

qlib/contrib/model/pytorch_gats.py

@@ -8,13 +8,9 @@
 import os
 import numpy as np
 import pandas as pd
+from typing import Text, Union
 import copy
-from ...utils import (
-    unpack_archive_with_buffer,
-    save_multiple_parts_file,
-    get_or_create_path,
-    drop_nan_by_y_index,
-)
+from ...utils import get_or_create_path
 from ...log import get_module_logger
 import torch
 import torch.nn as nn
@@ -83,7 +79,6 @@ def __init__(
         self.with_pretrain = with_pretrain
         self.model_path = model_path
         self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
-        self.use_gpu = torch.cuda.is_available()
         self.seed = seed
 
         self.logger.info(
@@ -310,11 +305,11 @@ def fit(
         if self.use_gpu:
             torch.cuda.empty_cache()
 
-    def predict(self, dataset):
+    def predict(self, dataset: DatasetH, segment: Union[Text, slice] = "test"):
         if not self.fitted:
             raise ValueError("model is not fitted yet!")
 
-        x_test = dataset.prepare("test", col_set="feature")
+        x_test = dataset.prepare(segment, col_set="feature")
         index = x_test.index
         self.GAT_model.eval()
         x_values = x_test.values

qlib/contrib/model/pytorch_gats_ts.py

@@ -9,12 +9,7 @@
 import numpy as np
 import pandas as pd
 import copy
-from ...utils import (
-    unpack_archive_with_buffer,
-    save_multiple_parts_file,
-    get_or_create_path,
-    drop_nan_by_y_index,
-)
+from ...utils import get_or_create_path
 from ...log import get_module_logger
 import torch
 import torch.nn as nn

qlib/contrib/model/pytorch_gru.py

@@ -8,13 +8,9 @@
 import os
 import numpy as np
 import pandas as pd
+from typing import Text, Union
 import copy
-from ...utils import (
-    unpack_archive_with_buffer,
-    save_multiple_parts_file,
-    get_or_create_path,
-    drop_nan_by_y_index,
-)
+from ...utils import get_or_create_path
 from ...log import get_module_logger
 
 import torch
@@ -273,11 +269,11 @@ def fit(
         if self.use_gpu:
             torch.cuda.empty_cache()
 
-    def predict(self, dataset):
+    def predict(self, dataset: DatasetH, segment: Union[Text, slice] = "test"):
         if not self.fitted:
             raise ValueError("model is not fitted yet!")
 
-        x_test = dataset.prepare("test", col_set="feature")
+        x_test = dataset.prepare(segment, col_set="feature", data_key=DataHandlerLP.DK_I)
         index = x_test.index
         self.gru_model.eval()
         x_values = x_test.values

qlib/contrib/model/pytorch_gru_ts.py

@@ -9,12 +9,7 @@
 import numpy as np
 import pandas as pd
 import copy
-from ...utils import (
-    unpack_archive_with_buffer,
-    save_multiple_parts_file,
-    get_or_create_path,
-    drop_nan_by_y_index,
-)
+from ...utils import get_or_create_path
 from ...log import get_module_logger
 
 import torch

qlib/contrib/model/pytorch_lstm.py

@@ -8,13 +8,9 @@
 import os
 import numpy as np
 import pandas as pd
+from typing import Text, Union
 import copy
-from ...utils import (
-    unpack_archive_with_buffer,
-    save_multiple_parts_file,
-    get_or_create_path,
-    drop_nan_by_y_index,
-)
+from ...utils import get_or_create_path
 from ...log import get_module_logger
 
 import torch
@@ -268,29 +264,25 @@ def fit(
         if self.use_gpu:
             torch.cuda.empty_cache()
 
-    def predict(self, dataset):
+    def predict(self, dataset: DatasetH, segment: Union[Text, slice] = "test"):
         if not self.fitted:
             raise ValueError("model is not fitted yet!")
 
-        x_test = dataset.prepare("test", col_set="feature")
+        x_test = dataset.prepare(segment, col_set="feature", data_key=DataHandlerLP.DK_I)
         index = x_test.index
         self.lstm_model.eval()
         x_values = x_test.values
         sample_num = x_values.shape[0]
         preds = []
 
         for begin in range(sample_num)[:: self.batch_size]:
-
             if sample_num - begin < self.batch_size:
                 end = sample_num
             else:
                 end = begin + self.batch_size
-
             x_batch = torch.from_numpy(x_values[begin:end]).float().to(self.device)
-
             with torch.no_grad():
                 pred = self.lstm_model(x_batch).detach().cpu().numpy()
-
             preds.append(pred)
 
         return pd.Series(np.concatenate(preds), index=index)

qlib/contrib/model/pytorch_lstm_ts.py

@@ -9,12 +9,7 @@
 import numpy as np
 import pandas as pd
 import copy
-from ...utils import (
-    unpack_archive_with_buffer,
-    save_multiple_parts_file,
-    get_or_create_path,
-    drop_nan_by_y_index,
-)
+from ...utils import get_or_create_path
 from ...log import get_module_logger
 
 import torch

qlib/contrib/model/pytorch_nn.py

@@ -8,6 +8,7 @@
 import os
 import numpy as np
 import pandas as pd
+from typing import Text, Union
 from sklearn.metrics import roc_auc_score, mean_squared_error
 
 import torch
@@ -18,7 +19,7 @@
 from ...model.base import Model
 from ...data.dataset import DatasetH
 from ...data.dataset.handler import DataHandlerLP
-from ...utils import unpack_archive_with_buffer, save_multiple_parts_file, get_or_create_path, drop_nan_by_y_index
+from ...utils import unpack_archive_with_buffer, save_multiple_parts_file, get_or_create_path
 from ...log import get_module_logger
 from ...workflow import R
 
@@ -48,8 +49,8 @@ class DNNModelPytorch(Model):
 
     def __init__(
         self,
-        input_dim,
-        output_dim,
+        input_dim=360,
+        output_dim=1,
         layers=(256,),
         lr=0.001,
         max_steps=300,
@@ -271,13 +272,12 @@ def get_loss(self, pred, w, target, loss_type):
         else:
             raise NotImplementedError("loss {} is not supported!".format(loss_type))
 
-    def predict(self, dataset):
+    def predict(self, dataset: DatasetH, segment: Union[Text, slice] = "test"):
         if not self.fitted:
             raise ValueError("model is not fitted yet!")
-        x_test_pd = dataset.prepare("test", col_set="feature")
+        x_test_pd = dataset.prepare(segment, col_set="feature", data_key=DataHandlerLP.DK_I)
         x_test = torch.from_numpy(x_test_pd.values).float().to(self.device)
         self.dnn_model.eval()
-
         with torch.no_grad():
             preds = self.dnn_model(x_test).detach().cpu().numpy()
         return pd.Series(np.squeeze(preds), index=x_test_pd.index)