iicd/xautodl
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update tests for torch/cuda

Browse Source
This commit is contained in:
D-X-Y
2021-03-30 12:05:52 +00:00
parent c2270fd153
commit e5ec43e04a
12 changed files with 386 additions and 248 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
lib/utils/flop_benchmark.py
View File

@@ -9,15 +9,19 @@ def count_parameters_in_MB(model):
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())
counts = sum(np.prod(v.size()) for v in model_or_parameters.parameters())
elif isinstance(models_or_parameters, nn.Parameter):
counts = models_or_parameters.numel()
elif isinstance(models_or_parameters, (list, tuple)):
counts = sum(count_parameters(x, None) for x in models_or_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
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
elif unit.lower() == "mb" or unit.lower() == "m":
counts /= 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
elif unit is not None:
raise ValueError("Unknow unit: {:}".format(unit))
return counts