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

10
qlib/contrib/model/pytorch_alstm.py
View File

@@ -23,6 +23,7 @@ import torch
import torch.nn as nn import torch.nn as nn
import torch.optim as optim import torch.optim as optim
from .pytorch_utils import count_parameters
from ...model.base import Model from ...model.base import Model
from ...data.dataset import DatasetH from ...data.dataset import DatasetH
from ...data.dataset.handler import DataHandlerLP from ...data.dataset.handler import DataHandlerLP
@@ -39,8 +40,8 @@ class ALSTM(Model):
the evaluate metric used in early stop the evaluate metric used in early stop
optimizer : str optimizer : str
optimizer name optimizer name
GPU : str GPU : int
the GPU ID(s) used for training the GPU ID used for training
""" """
def __init__( def __init__(
@@ -76,7 +77,7 @@ class ALSTM(Model):
self.early_stop = early_stop self.early_stop = early_stop
self.optimizer = optimizer.lower() self.optimizer = optimizer.lower()
self.loss = loss 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.use_gpu = torch.cuda.is_available()
self.seed = seed self.seed = seed
@@ -123,6 +124,9 @@ class ALSTM(Model):
num_layers=self.num_layers, num_layers=self.num_layers,
dropout=self.dropout, dropout=self.dropout,
) )
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": if optimizer.lower() == "adam":
self.train_optimizer = optim.Adam(self.ALSTM_model.parameters(), lr=self.lr) self.train_optimizer = optim.Adam(self.ALSTM_model.parameters(), lr=self.lr)
elif optimizer.lower() == "gd": elif optimizer.lower() == "gd":

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

@@ -40,8 +40,8 @@ class ALSTM(Model):
the evaluate metric used in early stop the evaluate metric used in early stop
optimizer : str optimizer : str
optimizer name optimizer name
GPU : str GPU : int
the GPU ID(s) used for training the GPU ID used for training
""" """
def __init__( def __init__(
@@ -78,7 +78,7 @@ class ALSTM(Model):
self.early_stop = early_stop self.early_stop = early_stop
self.optimizer = optimizer.lower() self.optimizer = optimizer.lower()
self.loss = loss 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.n_jobs = n_jobs self.n_jobs = n_jobs
self.use_gpu = torch.cuda.is_available() self.use_gpu = torch.cuda.is_available()
self.seed = seed self.seed = seed

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

@@ -23,6 +23,7 @@ import torch
import torch.nn as nn import torch.nn as nn
import torch.optim as optim import torch.optim as optim
from .pytorch_utils import count_parameters
from ...model.base import Model from ...model.base import Model
from ...data.dataset import DatasetH from ...data.dataset import DatasetH
from ...data.dataset.handler import DataHandlerLP from ...data.dataset.handler import DataHandlerLP
@@ -76,7 +77,7 @@ class GRU(Model):
self.early_stop = early_stop self.early_stop = early_stop
self.optimizer = optimizer.lower() self.optimizer = optimizer.lower()
self.loss = loss 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.use_gpu = torch.cuda.is_available()
self.seed = seed self.seed = seed
@@ -123,6 +124,9 @@ class GRU(Model):
num_layers=self.num_layers, num_layers=self.num_layers,
dropout=self.dropout, dropout=self.dropout,
) )
self.logger.info("model:\n{:}".format(self.gru_model))
self.logger.info("model size: {:.4f} MB".format(count_parameters(self.gru_model)))
if optimizer.lower() == "adam": if optimizer.lower() == "adam":
self.train_optimizer = optim.Adam(self.gru_model.parameters(), lr=self.lr) self.train_optimizer = optim.Adam(self.gru_model.parameters(), lr=self.lr)
elif optimizer.lower() == "gd": elif optimizer.lower() == "gd":

8
qlib/contrib/model/pytorch_gru_ts.py
View File

@@ -24,6 +24,7 @@ import torch.nn as nn
import torch.optim as optim import torch.optim as optim
from torch.utils.data import DataLoader from torch.utils.data import DataLoader
from .pytorch_utils import count_parameters
from ...model.base import Model from ...model.base import Model
from ...data.dataset import DatasetH, TSDatasetH from ...data.dataset import DatasetH, TSDatasetH
from ...data.dataset.handler import DataHandlerLP from ...data.dataset.handler import DataHandlerLP
@@ -78,7 +79,7 @@ class GRU(Model):
self.early_stop = early_stop self.early_stop = early_stop
self.optimizer = optimizer.lower() self.optimizer = optimizer.lower()
self.loss = loss 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.n_jobs = n_jobs self.n_jobs = n_jobs
self.use_gpu = torch.cuda.is_available() self.use_gpu = torch.cuda.is_available()
self.seed = seed self.seed = seed
@@ -127,7 +128,10 @@ class GRU(Model):
hidden_size=self.hidden_size, hidden_size=self.hidden_size,
num_layers=self.num_layers, num_layers=self.num_layers,
dropout=self.dropout, dropout=self.dropout,
).to(self.device) )
self.logger.info("model:\n{:}".format(self.gru_model))
self.logger.info("model size: {:.4f} MB".format(count_parameters(self.gru_model)))
if optimizer.lower() == "adam": if optimizer.lower() == "adam":
self.train_optimizer = optim.Adam(self.GRU_model.parameters(), lr=self.lr) self.train_optimizer = optim.Adam(self.GRU_model.parameters(), lr=self.lr)
elif optimizer.lower() == "gd": elif optimizer.lower() == "gd":

2
qlib/contrib/model/pytorch_lstm.py
View File

@@ -76,7 +76,7 @@ class LSTM(Model):
self.early_stop = early_stop self.early_stop = early_stop
self.optimizer = optimizer.lower() self.optimizer = optimizer.lower()
self.loss = loss 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.use_gpu = torch.cuda.is_available()
self.seed = seed self.seed = seed

2
qlib/contrib/model/pytorch_lstm_ts.py
View File

@@ -78,7 +78,7 @@ class LSTM(Model):
self.early_stop = early_stop self.early_stop = early_stop
self.optimizer = optimizer.lower() self.optimizer = optimizer.lower()
self.loss = loss 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.n_jobs = n_jobs self.n_jobs = n_jobs
self.use_gpu = torch.cuda.is_available() self.use_gpu = torch.cuda.is_available()
self.seed = seed self.seed = seed

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

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

21
qlib/contrib/model/pytorch_utils.py Normal file
View File

@@ -0,0 +1,21 @@
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
import numpy as np
import torch.nn as nn
def count_parameters(model_or_parameters, unit="mb"):
if isinstance(model_or_parameters, nn.Module):
counts = np.sum(np.prod(v.size()) for v in model_or_parameters.parameters())
else:
counts = np.sum(np.prod(v.size()) for v in model_or_parameters)
if unit.lower() == "mb":
counts /= 1e6
elif unit.lower() == "kb":
counts /= 1e3
elif unit.lower() == "gb":
counts /= 1e9
elif unit is not None:
raise ValueError("Unknow unit: {:}".format(unit))
return counts