Show model size for pytorch models

qlib/contrib/model/pytorch_alstm_ts.py

@@ -24,6 +24,7 @@ import torch.nn as nn
 import torch.optim as optim
 from torch.utils.data import DataLoader
 
+from .pytorch_utils import count_parameters
 from ...model.base import Model
 from ...data.dataset import DatasetH, TSDatasetH
 from ...data.dataset.handler import DataHandlerLP
@@ -127,7 +128,10 @@ class ALSTM(Model):
             hidden_size=self.hidden_size,
             num_layers=self.num_layers,
             dropout=self.dropout,
-        ).to(self.device)
+        )
+        self.logger.info("model:\n{:}".format(self.ALSTM_model))
+        self.logger.info("model size: {:.4f} MB".format(count_parameters(self.ALSTM_model)))
+
         if optimizer.lower() == "adam":
             self.train_optimizer = optim.Adam(self.ALSTM_model.parameters(), lr=self.lr)
         elif optimizer.lower() == "gd":

qlib/contrib/model/pytorch_gats.py

@@ -22,6 +22,7 @@ import torch
 import torch.nn as nn
 import torch.optim as optim
 
+from .pytorch_utils import count_parameters
 from ...model.base import Model
 from ...data.dataset import DatasetH
 from ...data.dataset.handler import DataHandlerLP
@@ -42,8 +43,8 @@ class GATs(Model):
         the evaluate metric used in early stop
     optimizer : str
         optimizer name
-    GPU : str
+    GPU : int
-        the GPU ID(s) used for training
+        the GPU ID used for training
     """
 
     def __init__(
@@ -83,7 +84,7 @@ class GATs(Model):
         self.base_model = base_model
         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.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
 
@@ -135,6 +136,9 @@ class GATs(Model):
             dropout=self.dropout,
             base_model=self.base_model,
         )
+        self.logger.info("model:\n{:}".format(self.GAT_model))
+        self.logger.info("model size: {:.4f} MB".format(count_parameters(self.GAT_model)))
+
         if optimizer.lower() == "adam":
             self.train_optimizer = optim.Adam(self.GAT_model.parameters(), lr=self.lr)
         elif optimizer.lower() == "gd":

qlib/contrib/model/pytorch_gats_ts.py

@@ -24,6 +24,7 @@ import torch.optim as optim
 from torch.utils.data import DataLoader
 from torch.utils.data import Sampler
 
+from .pytorch_utils import count_parameters
 from ...model.base import Model
 from ...data.dataset import DatasetH
 from ...data.dataset.handler import DataHandlerLP
@@ -62,8 +63,8 @@ class GATs(Model):
         the evaluate metric used in early stop
     optimizer : str
         optimizer name
-    GPU : str
+    GPU : int
-        the GPU ID(s) used for training
+        the GPU ID used for training
     """
 
     def __init__(
@@ -104,7 +105,7 @@ class GATs(Model):
         self.base_model = base_model
         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.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
         self.n_jobs = n_jobs
         self.use_gpu = torch.cuda.is_available()
         self.seed = seed
@@ -157,6 +158,9 @@ class GATs(Model):
             dropout=self.dropout,
             base_model=self.base_model,
         )
+        self.logger.info("model:\n{:}".format(self.GAT_model))
+        self.logger.info("model size: {:.4f} MB".format(count_parameters(self.GAT_model)))
+
         if optimizer.lower() == "adam":
             self.train_optimizer = optim.Adam(self.GAT_model.parameters(), lr=self.lr)
         elif optimizer.lower() == "gd":

qlib/contrib/model/pytorch_nn.py

@@ -15,6 +15,7 @@ import torch
 import torch.nn as nn
 import torch.optim as optim
 
+from .pytorch_utils import count_parameters
 from ...model.base import Model
 from ...data.dataset import DatasetH
 from ...data.dataset.handler import DataHandlerLP
@@ -129,6 +130,9 @@ class DNNModelPytorch(Model):
         self._scorer = mean_squared_error if loss == "mse" else roc_auc_score
 
         self.dnn_model = Net(input_dim, output_dim, layers, loss=self.loss_type)
+        self.logger.info("model:\n{:}".format(self.dnn_model))
+        self.logger.info("model size: {:.4f} MB".format(count_parameters(self.dnn_model)))
+
         if optimizer.lower() == "adam":
             self.train_optimizer = optim.Adam(self.dnn_model.parameters(), lr=self.lr, weight_decay=self.weight_decay)
         elif optimizer.lower() == "gd":

qlib/contrib/model/pytorch_sfm.py

@@ -23,6 +23,7 @@ import torch.nn as nn
 import torch.nn.init as init
 import torch.optim as optim
 
+from .pytorch_utils import count_parameters
 from ...model.base import Model
 from ...data.dataset import DatasetH
 from ...data.dataset.handler import DataHandlerLP
@@ -196,8 +197,8 @@ class SFM(Model):
         learning rate
     optimizer : str
         optimizer name
-    GPU : str
+    GPU : int
-        the GPU ID(s) used for training
+        the GPU ID used for training
     """
 
     def __init__(
@@ -216,7 +217,7 @@ class SFM(Model):
         eval_steps=5,
         loss="mse",
         optimizer="gd",
-        GPU="0",
+        GPU=0,
         seed=None,
         **kwargs
     ):
@@ -239,7 +240,7 @@ class SFM(Model):
         self.eval_steps = eval_steps
         self.optimizer = optimizer.lower()
         self.loss = loss
-        self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() else "cpu")
+        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
 
@@ -295,6 +296,9 @@ class SFM(Model):
             dropout_U=self.dropout_U,
             device=self.device,
         )
+        self.logger.info("model:\n{:}".format(self.sfm_model))
+        self.logger.info("model size: {:.4f} MB".format(count_parameters(self.sfm_model)))
+
         if optimizer.lower() == "adam":
             self.train_optimizer = optim.Adam(self.sfm_model.parameters(), lr=self.lr)
         elif optimizer.lower() == "gd":

qlib/contrib/model/pytorch_tabnet.py

@@ -23,6 +23,7 @@ import torch.optim as optim
 import torch.nn.functional as F
 from torch.autograd import Function
 
+from .pytorch_utils import count_parameters
 from ...model.base import Model
 from ...data.dataset import DatasetH
 from ...data.dataset.handler import DataHandlerLP
@@ -49,7 +50,7 @@ class TabnetModel(Model):
         loss="mse",
         metric="",
         early_stop=20,
-        GPU="1",
+        GPU=0,
         pretrain_loss="custom",
         ps=0.3,
         lr=0.01,
@@ -75,7 +76,7 @@ class TabnetModel(Model):
         self.n_epochs = n_epochs
         self.logger = get_module_logger("TabNet")
         self.pretrain_n_epochs = pretrain_n_epochs
-        self.device = "cuda:%s" % (GPU) if torch.cuda.is_available() else "cpu"
+        self.device = "cuda:%s" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu"
         self.loss = loss
         self.metric = metric
         self.early_stop = early_stop
@@ -98,6 +99,8 @@ class TabnetModel(Model):
         self.tabnet_decoder = TabNet_Decoder(self.out_dim, self.d_feat, n_shared, n_ind, vbs, n_steps, self.device).to(
             self.device
         )
+        self.logger.info("model:\n{:}\n{:}".format(self.tabnet_model, self.tabnet_decoder))
+        self.logger.info("model size: {:.4f} MB".format(count_parameters(self.tabnet_model) + count_parameters(self.tabnet_decoder)))
 
         if optimizer.lower() == "adam":
             self.pretrain_optimizer = optim.Adam(