Update models.

HugoZou · Dec 9, 2020 · ec0d783 · ec0d783
1 parent 752f17e
commit ec0d783
Show file tree

Hide file tree

Showing 5 changed files with 248 additions and 206 deletions.
diff --git a/examples/benchmarks/GRU/csi300_gru_ts.pkl b/examples/benchmarks/GRU/csi300_gru_ts.pkl
diff --git a/examples/benchmarks/LSTM/csi300_lstm_ts.pkl b/examples/benchmarks/LSTM/csi300_lstm_ts.pkl
diff --git a/...s/SFM/workflow_config_alstm_Alpha158.yaml → ...rks/SFM/workflow_config_sfm_Alpha158.yaml b/...s/SFM/workflow_config_alstm_Alpha158.yaml → ...rks/SFM/workflow_config_sfm_Alpha158.yaml
@@ -57,7 +57,7 @@ task:
             num_layers: 2
             dropout: 0.0
             n_epochs: 200
-            lr: 1e-3
+            lr: 1e-2
             early_stop: 10
             batch_size: 800
             metric: loss

diff --git a/qlib/contrib/model/pytorch_gats_ts.py b/qlib/contrib/model/pytorch_gats_ts.py
@@ -21,6 +21,7 @@
 import torch
 import torch.nn as nn
 import torch.optim as optim
+from torch.utils.data import DataLoader
 
 from ...model.base import Model
 from ...data.dataset import DatasetH
@@ -62,7 +63,9 @@ def __init__(
         model_path=None,
         optimizer="adam",
         GPU="0",
+        n_jobs=10,
         seed=None,
+        batch_size=800,
         **kwargs
     ):
         # Set logger.
@@ -84,8 +87,10 @@ def __init__(
         self.with_pretrain = with_pretrain
         self.model_path = model_path
         self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() else "cpu")
+        self.n_jobs = n_jobs
         self.use_gpu = torch.cuda.is_available()
         self.seed = seed
+        self.batch_size = batch_size
 
         self.logger.info(
             "GATs parameters setting:"
@@ -218,7 +223,7 @@ def test_epoch(self, data_loader):
 
     def fit(
         self,
-        dataset: DatasetH,
+        dataset,
         evals_result=dict(),
         verbose=True,
         save_path=None,