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

Remove set_log_basic_config, refine count_parameters, rename root_uri as get_local_dir

Browse Source
This commit is contained in:
D-X-Y
2021-03-11 02:33:00 +00:00
parent e061443560
commit f6ed175070
4 changed files with 33 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
qlib/contrib/model/pytorch_utils.py
View File

@@ -4,7 +4,19 @@
import torch.nn as nn
def count_parameters(models_or_parameters, unit="mb"):
def count_parameters(models_or_parameters, unit="m"):
"""
This function is to obtain the storage size unit of a (or multiple) models.
Parameters
----------
models_or_parameters : PyTorch model(s) or a list of parameters.
unit : the storage size unit.
Returns
-------
The number of parameters of the given model(s) or parameters.
"""
if isinstance(models_or_parameters, nn.Module):
counts = sum(v.numel() for v in models_or_parameters.parameters())
elif isinstance(models_or_parameters, nn.Parameter):
@@ -13,12 +25,13 @@ def count_parameters(models_or_parameters, unit="mb"):
return sum(count_parameters(x, unit) for x in models_or_parameters)
else:
counts = sum(v.numel() for v in models_or_parameters)
if unit.lower() == "mb":
counts /= 1e6
elif unit.lower() == "kb":
counts /= 1e3
elif unit.lower() == "gb":
counts /= 1e9
unit = unit.lower()
if unit == "kb" or unit == "k":
counts /= 2 ** 10
elif unit == "mb" or unit == "m":
counts /= 2 ** 20
elif unit == "gb" or unit == "g":
counts /= 2 ** 30
elif unit is not None:
raise ValueError("Unknow unit: {:}".format(unit))
return counts