liuyu/qlib
1
0
Fork 0
mirror of https://github.com/microsoft/qlib.git synced 2026-07-14 16:26:55 +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

4
qlib/contrib/model/pytorch_nn.py
View File

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