diff --git a/qlib/contrib/model/pytorch_alstm_ts.py b/qlib/contrib/model/pytorch_alstm_ts.py index c8854e8d3..933905439 100644 --- a/qlib/contrib/model/pytorch_alstm_ts.py +++ b/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": diff --git a/qlib/contrib/model/pytorch_gats.py b/qlib/contrib/model/pytorch_gats.py index 07048e1bc..061feaa84 100644 --- a/qlib/contrib/model/pytorch_gats.py +++ b/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 - the GPU ID(s) used for training + GPU : int + 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": diff --git a/qlib/contrib/model/pytorch_gats_ts.py b/qlib/contrib/model/pytorch_gats_ts.py index 1e94f56e4..b179a6df6 100644 --- a/qlib/contrib/model/pytorch_gats_ts.py +++ b/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 - the GPU ID(s) used for training + GPU : int + 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": diff --git a/qlib/contrib/model/pytorch_nn.py b/qlib/contrib/model/pytorch_nn.py index f8b0b7748..a51481c85 100644 --- a/qlib/contrib/model/pytorch_nn.py +++ b/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": diff --git a/qlib/contrib/model/pytorch_sfm.py b/qlib/contrib/model/pytorch_sfm.py index d5169e6c7..55058797a 100644 --- a/qlib/contrib/model/pytorch_sfm.py +++ b/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 - the GPU ID(s) used for training + GPU : int + 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": diff --git a/qlib/contrib/model/pytorch_tabnet.py b/qlib/contrib/model/pytorch_tabnet.py index 62e32d701..ff56b0561 100644 --- a/qlib/contrib/model/pytorch_tabnet.py +++ b/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(