Update all baseline models.

qlib/contrib/model/pytorch_sfm.py

@@ -19,7 +19,12 @@ import pandas as pd
 import copy
 from sklearn.metrics import roc_auc_score, mean_squared_error
 import logging
-from ...utils import unpack_archive_with_buffer, save_multiple_parts_file, create_save_path, drop_nan_by_y_index
+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
@@ -33,7 +38,16 @@ from ...data.dataset.handler import DataHandlerLP
 
 
 class SFM_Model(nn.Module):
-    def __init__(self, d_feat=6, output_dim=1, freq_dim=10, hidden_size=64, dropout_W=0.0, dropout_U=0.0, device="cpu"):
+    def __init__(
+        self,
+        d_feat=6,
+        output_dim=1,
+        freq_dim=10,
+        hidden_size=64,
+        dropout_W=0.0,
+        dropout_U=0.0,
+        device="cpu",
+    ):
         super().__init__()
 
         self.input_dim = d_feat
@@ -42,30 +56,52 @@ class SFM_Model(nn.Module):
         self.hidden_dim = hidden_size
         self.device = device
 
-        self.W_i = nn.Parameter(init.xavier_uniform_(torch.empty((self.input_dim, self.hidden_dim))))
-        self.U_i = nn.Parameter(init.orthogonal_(torch.empty(self.hidden_dim, self.hidden_dim)))
+        self.W_i = nn.Parameter(
+            init.xavier_uniform_(torch.empty((self.input_dim, self.hidden_dim)))
+        )
+        self.U_i = nn.Parameter(
+            init.orthogonal_(torch.empty(self.hidden_dim, self.hidden_dim))
+        )
         self.b_i = nn.Parameter(torch.zeros(self.hidden_dim))
 
-        self.W_ste = nn.Parameter(init.xavier_uniform_(torch.empty(self.input_dim, self.hidden_dim)))
-        self.U_ste = nn.Parameter(init.orthogonal_(torch.empty(self.hidden_dim, self.hidden_dim)))
+        self.W_ste = nn.Parameter(
+            init.xavier_uniform_(torch.empty(self.input_dim, self.hidden_dim))
+        )
+        self.U_ste = nn.Parameter(
+            init.orthogonal_(torch.empty(self.hidden_dim, self.hidden_dim))
+        )
         self.b_ste = nn.Parameter(torch.ones(self.hidden_dim))
 
-        self.W_fre = nn.Parameter(init.xavier_uniform_(torch.empty(self.input_dim, self.freq_dim)))
-        self.U_fre = nn.Parameter(init.orthogonal_(torch.empty(self.hidden_dim, self.freq_dim)))
+        self.W_fre = nn.Parameter(
+            init.xavier_uniform_(torch.empty(self.input_dim, self.freq_dim))
+        )
+        self.U_fre = nn.Parameter(
+            init.orthogonal_(torch.empty(self.hidden_dim, self.freq_dim))
+        )
         self.b_fre = nn.Parameter(torch.ones(self.freq_dim))
 
-        self.W_c = nn.Parameter(init.xavier_uniform_(torch.empty(self.input_dim, self.hidden_dim)))
-        self.U_c = nn.Parameter(init.orthogonal_(torch.empty(self.hidden_dim, self.hidden_dim)))
+        self.W_c = nn.Parameter(
+            init.xavier_uniform_(torch.empty(self.input_dim, self.hidden_dim))
+        )
+        self.U_c = nn.Parameter(
+            init.orthogonal_(torch.empty(self.hidden_dim, self.hidden_dim))
+        )
         self.b_c = nn.Parameter(torch.zeros(self.hidden_dim))
 
-        self.W_o = nn.Parameter(init.xavier_uniform_(torch.empty(self.input_dim, self.hidden_dim)))
-        self.U_o = nn.Parameter(init.orthogonal_(torch.empty(self.hidden_dim, self.hidden_dim)))
+        self.W_o = nn.Parameter(
+            init.xavier_uniform_(torch.empty(self.input_dim, self.hidden_dim))
+        )
+        self.U_o = nn.Parameter(
+            init.orthogonal_(torch.empty(self.hidden_dim, self.hidden_dim))
+        )
         self.b_o = nn.Parameter(torch.zeros(self.hidden_dim))
 
         self.U_a = nn.Parameter(init.orthogonal_(torch.empty(self.freq_dim, 1)))
         self.b_a = nn.Parameter(torch.zeros(self.hidden_dim))
 
-        self.W_p = nn.Parameter(init.xavier_uniform_(torch.empty(self.hidden_dim, self.output_dim)))
+        self.W_p = nn.Parameter(
+            init.xavier_uniform_(torch.empty(self.hidden_dim, self.output_dim))
+        )
         self.b_p = nn.Parameter(torch.zeros(self.output_dim))
 
         self.activation = nn.Tanh()
@@ -101,8 +137,12 @@ class SFM_Model(nn.Module):
             x_o = torch.matmul(x * B_W[0], self.W_o) + self.b_o
 
             i = self.inner_activation(x_i + torch.matmul(h_tm1 * B_U[0], self.U_i))
-            ste = self.inner_activation(x_ste + torch.matmul(h_tm1 * B_U[0], self.U_ste))
-            fre = self.inner_activation(x_fre + torch.matmul(h_tm1 * B_U[0], self.U_fre))
+            ste = self.inner_activation(
+                x_ste + torch.matmul(h_tm1 * B_U[0], self.U_ste)
+            )
+            fre = self.inner_activation(
+                x_fre + torch.matmul(h_tm1 * B_U[0], self.U_fre)
+            )
 
             ste = torch.reshape(ste, (-1, self.hidden_dim, 1))
             fre = torch.reshape(fre, (-1, 1, self.freq_dim))
@@ -157,7 +197,16 @@ class SFM_Model(nn.Module):
 
         init_state_time = torch.tensor(0).to(self.device)
 
-        self.states = [init_state_p, init_state_h, init_state_S_re, init_state_S_im, init_state_time, None, None, None]
+        self.states = [
+            init_state_p,
+            init_state_h,
+            init_state_S_re,
+            init_state_S_im,
+            init_state_time,
+            None,
+            None,
+            None,
+        ]
 
     def get_constants(self, x):
         constants = []
@@ -282,7 +331,9 @@ class SFM(Model):
         elif optimizer.lower() == "gd":
             self.train_optimizer = optim.SGD(self.sfm_model.parameters(), lr=self.lr)
         else:
-            raise NotImplementedError("optimizer {} is not supported!".format(optimizer))
+            raise NotImplementedError(
+                "optimizer {} is not supported!".format(optimizer)
+            )
 
         self._fitted = False
         self.sfm_model.to(self.device)
@@ -305,8 +356,16 @@ class SFM(Model):
             if len(indices) - i < self.batch_size:
                 break
 
-            feature = torch.from_numpy(x_values[indices[i : i + self.batch_size]]).float().to(self.device)
-            label = torch.from_numpy(y_values[indices[i : i + self.batch_size]]).float().to(self.device)
+            feature = (
+                torch.from_numpy(x_values[indices[i : i + self.batch_size]])
+                .float()
+                .to(self.device)
+            )
+            label = (
+                torch.from_numpy(y_values[indices[i : i + self.batch_size]])
+                .float()
+                .to(self.device)
+            )
 
             pred = self.sfm_model(feature)
             loss = self.loss_fn(pred, label)
@@ -332,8 +391,16 @@ class SFM(Model):
             if len(indices) - i < self.batch_size:
                 break
 
-            feature = torch.from_numpy(x_train_values[indices[i : i + self.batch_size]]).float().to(self.device)
-            label = torch.from_numpy(y_train_values[indices[i : i + self.batch_size]]).float().to(self.device)
+            feature = (
+                torch.from_numpy(x_train_values[indices[i : i + self.batch_size]])
+                .float()
+                .to(self.device)
+            )
+            label = (
+                torch.from_numpy(y_train_values[indices[i : i + self.batch_size]])
+                .float()
+                .to(self.device)
+            )
 
             pred = self.sfm_model(feature)
             loss = self.loss_fn(pred, label)
@@ -352,7 +419,9 @@ class SFM(Model):
     ):
 
         df_train, df_valid = dataset.prepare(
-            ["train", "valid"], col_set=["feature", "label"], data_key=DataHandlerLP.DK_L
+            ["train", "valid"],
+            col_set=["feature", "label"],
+            data_key=DataHandlerLP.DK_L,
         )
         x_train, y_train = df_train["feature"], df_train["label"]
         x_valid, y_valid = df_valid["feature"], df_valid["label"]
@@ -409,7 +478,7 @@ class SFM(Model):
 
         mask = torch.isfinite(label)
 
-        if self.metric == "" or self.metric == "loss":  # use loss
+        if self.metric == "" or self.metric == "loss":
             return -self.loss_fn(pred[mask], label[mask])
 
         raise ValueError("unknown metric `%s`" % self.metric)