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

Add count_parameters for pytorch models in contrib

Browse Source
This commit is contained in:
D-X-Y
2021-03-05 12:07:43 +00:00
parent b14a559a52
commit 131f0e2e67
8 changed files with 47 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -42,8 +42,8 @@ class DNNModelPytorch(Model):
learning rate decay steps
optimizer : str
optimizer name
GPU : str
the GPU ID(s) used for training
GPU : int
the GPU ID used for training
"""
def __init__(
@@ -80,7 +80,7 @@ class DNNModelPytorch(Model):
self.lr_decay_steps = lr_decay_steps
self.optimizer = optimizer.lower()
self.loss_type = 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
self.weight_decay = weight_decay