liuyu/qlib
1
0
Fork 0
mirror of https://github.com/microsoft/qlib.git synced 2026-07-15 16:56:54 +08:00
Code Activity

Show model size for pytorch models

Browse Source
This commit is contained in:
D-X-Y
2021-03-05 12:46:41 +00:00
parent 131f0e2e67
commit 49697b1f15
6 changed files with 36 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
qlib/contrib/model/pytorch_sfm.py
View File

@@ -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":