This commit is contained in:
Glenn Jocher 2019-03-19 15:50:15 +02:00
parent 973715060d
commit 613ce1be1a
1 changed files with 2 additions and 2 deletions

View File

@ -37,10 +37,10 @@ def model_info(model):
# Plots a line-by-line description of a PyTorch model # Plots a line-by-line description of a PyTorch model
n_p = sum(x.numel() for x in model.parameters()) # number parameters n_p = sum(x.numel() for x in model.parameters()) # number parameters
n_g = sum(x.numel() for x in model.parameters() if x.requires_grad) # number gradients n_g = sum(x.numel() for x in model.parameters() if x.requires_grad) # number gradients
print('\n%5s %35s %9s %12s %20s %12s %12s' % ('layer', 'name', 'gradient', 'parameters', 'shape', 'mu', 'sigma')) print('\n%5s %38s %9s %12s %20s %12s %12s' % ('layer', 'name', 'gradient', 'parameters', 'shape', 'mu', 'sigma'))
for i, (name, p) in enumerate(model.named_parameters()): for i, (name, p) in enumerate(model.named_parameters()):
name = name.replace('module_list.', '') name = name.replace('module_list.', '')
print('%5g %35s %9s %12g %20s %12.3g %12.3g' % ( print('%5g %38s %9s %12g %20s %12.3g %12.3g' % (
i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std())) i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std()))
print('Model Summary: %g layers, %g parameters, %g gradients' % (i + 1, n_p, n_g)) print('Model Summary: %g layers, %g parameters, %g gradients' % (i + 1, n_p, n_g))