This commit is contained in:
D-X-Y
2021-05-26 04:47:38 +00:00
parent 5eab0de53e
commit d557c328a8
4 changed files with 23 additions and 27 deletions

View File

@@ -1,6 +1,3 @@
#####################################################
# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.03 #
#####################################################
import math
import abc
import numpy as np
@@ -13,11 +10,11 @@ class UnifiedSplit:
"""A class to unify the split strategy."""
def __init__(self, total_num, mode):
# Training Set 65%
num_of_train = int(total_num * 0.65)
# Training Set 75%
num_of_train = int(total_num * 0.75)
# Validation Set 05%
num_of_valid = int(total_num * 0.05)
# Test Set 30%
# Test Set 20%
num_of_set = total_num - num_of_train - num_of_valid
all_indexes = list(range(total_num))
if mode is None:

View File

@@ -1,27 +1,32 @@
#####################################################
# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.01 #
#####################################################
import torch
import torch.nn as nn
import numpy as np
def count_parameters_in_MB(model):
return count_parameters(model, "mb")
return count_parameters(model, "mb", deprecated=True)
def count_parameters(model_or_parameters, unit="mb"):
def count_parameters(model_or_parameters, unit="mb", deprecated=False):
if isinstance(model_or_parameters, nn.Module):
counts = sum(np.prod(v.size()) for v in model_or_parameters.parameters())
elif isinstance(model_or_parameters, nn.Parameter):
counts = models_or_parameters.numel()
elif isinstance(model_or_parameters, (list, tuple)):
counts = sum(count_parameters(x, None) for x in models_or_parameters)
counts = sum(
count_parameters(x, None, deprecated) for x in models_or_parameters
)
else:
counts = sum(np.prod(v.size()) for v in model_or_parameters)
if unit.lower() == "kb" or unit.lower() == "k":
counts /= 2 ** 10 # changed from 1e3 to 2^10
counts /= 1e3 if deprecated else 2 ** 10 # changed from 1e3 to 2^10
elif unit.lower() == "mb" or unit.lower() == "m":
counts /= 2 ** 20 # changed from 1e6 to 2^20
counts /= 1e6 if deprecated else 2 ** 20 # changed from 1e6 to 2^20
elif unit.lower() == "gb" or unit.lower() == "g":
counts /= 2 ** 30 # changed from 1e9 to 2^30
counts /= 1e9 if deprecated else 2 ** 30 # changed from 1e9 to 2^30
elif unit is not None:
raise ValueError("Unknow unit: {:}".format(unit))
return counts