mirror of
https://github.com/wassname/torchsummaryX.git
synced 2026-06-27 16:32:27 +08:00
Merge pull request #2 from wassname/master
Use module path names, use pandas
This commit is contained in:
Namhyuk Ahn
@@ -36,18 +36,19 @@ class Net(nn.Module):
|
||||
summary(Net(), torch.zeros((1, 1, 28, 28)))
|
||||
```
|
||||
```
|
||||
----------------------------------------------------------------------------------------------------
|
||||
Layer Kernel Shape Output Shape # Params (K) # Mult-Adds (M)
|
||||
====================================================================================================
|
||||
0_Conv2d [1, 10, 5, 5] [1, 10, 24, 24] 0.26 0.14
|
||||
1_Conv2d [10, 20, 5, 5] [1, 20, 8, 8] 5.02 0.32
|
||||
2_Dropout2d - [1, 20, 8, 8] - -
|
||||
3_Linear [320, 50] [1, 50] 16.05 0.02
|
||||
4_Linear [50, 10] [1, 10] 0.51 0.00
|
||||
Kernel Shape Output Shape Params (K) Mult-Adds (M)
|
||||
Layer
|
||||
0_conv1 [1, 10, 5, 5] [1, 10, 24, 24] 0.26 0.1440
|
||||
1_conv2 [10, 20, 5, 5] [1, 20, 8, 8] 5.02 0.3200
|
||||
2_conv2_drop - [1, 20, 8, 8] NaN NaN
|
||||
3_fc1 [320, 50] [1, 50] 16.05 0.0160
|
||||
4_fc2 [50, 10] [1, 10] 0.51 0.0005
|
||||
====================================================================================================
|
||||
Params (K) 21.8400
|
||||
Mult-Adds (M) 0.4805
|
||||
dtype: float64
|
||||
====================================================================================================
|
||||
# Params: 21.84K
|
||||
# Mult-Adds: 0.48M
|
||||
----------------------------------------------------------------------------------------------------
|
||||
```
|
||||
|
||||
RNN
|
||||
@@ -74,20 +75,18 @@ inputs = torch.zeros((100, 1), dtype=torch.long) # [length, batch_size]
|
||||
summary(Net(), inputs)
|
||||
```
|
||||
```
|
||||
----------------------------------------------------------------------------------------------------
|
||||
Layer Kernel Shape Output Shape # Params (K) # Mult-Adds (M)
|
||||
====================================================================================================
|
||||
0_Embedding [300, 20] [100, 1, 300] 6.00 0.01
|
||||
1_LSTM - [100, 1, 512] 3,768.32 3.76
|
||||
weight_ih_l0 [2048, 300]
|
||||
weight_hh_l0 [2048, 512]
|
||||
weight_ih_l1 [2048, 512]
|
||||
weight_hh_l1 [2048, 512]
|
||||
2_Linear [512, 20] [100, 1, 20] 10.26 0.01
|
||||
Kernel Shape Output Shape Params (K) Mult-Adds (M)
|
||||
Layer
|
||||
0_embedding [300, 20] [100, 1, 300] 6.00 0.006000
|
||||
1_encoder - [100, 1, 512] 3768.32 3.760128
|
||||
2_decoder [512, 20] [100, 1, 20] 10.26 0.010240
|
||||
====================================================================================================
|
||||
# Params: 3,784.58K
|
||||
# Mult-Adds: 3.78M
|
||||
----------------------------------------------------------------------------------------------------
|
||||
Params (K) 3784.580000
|
||||
Mult-Adds (M) 3.776368
|
||||
dtype: float64
|
||||
====================================================================================================
|
||||
|
||||
```
|
||||
|
||||
Recursive NN
|
||||
@@ -107,8 +106,16 @@ summary(Net(), torch.zeros((1, 64, 28, 28)))
|
||||
----------------------------------------------------------------------------------------------------
|
||||
Layer Kernel Shape Output Shape # Params (K) # Mult-Adds (M)
|
||||
====================================================================================================
|
||||
0_Conv2d [64, 64, 3, 3] [1, 64, 28, 28] 36.93 28.90
|
||||
1_Conv2d [64, 64, 3, 3] [1, 64, 28, 28] (recursive) 28.90
|
||||
Kernel Shape Output Shape Params (K) Mult-Adds (M)
|
||||
Layer
|
||||
0_conv1 [64, 64, 3, 3] [1, 64, 28, 28] 36.928 28.901376
|
||||
1_conv1 [64, 64, 3, 3] [1, 64, 28, 28] NaN 28.901376
|
||||
====================================================================================================
|
||||
Kernel Shape [64, 64, 3, 3, 64, 64, 3, 3]
|
||||
Output Shape [1, 64, 28, 28, 1, 64, 28, 28]
|
||||
Params (K) 36.928
|
||||
Mult-Adds (M) 57.8028
|
||||
dtype: object
|
||||
====================================================================================================
|
||||
# Params: 36.93K
|
||||
# Mult-Adds: 57.80M
|
||||
|
||||
@@ -1,12 +1,12 @@
|
||||
from collections import OrderedDict
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import torch
|
||||
|
||||
def summary(model, x, *args, **kwargs):
|
||||
"""Summarize the given input model.
|
||||
Summarized information are 1) output shape, 2) kernel shape,
|
||||
3) number of the parameters and 4) operations (Mult-Adds)
|
||||
|
||||
Args:
|
||||
model (Module): Model to summarize
|
||||
x (Tensor): Input tensor of the model with [N, C, H, W] shape
|
||||
@@ -17,7 +17,11 @@ def summary(model, x, *args, **kwargs):
|
||||
def hook(module, inputs, outputs):
|
||||
cls_name = str(module.__class__).split(".")[-1].split("'")[0]
|
||||
module_idx = len(summary)
|
||||
key = "{}_{}".format(module_idx, cls_name)
|
||||
|
||||
# Lookup name in a dict that includes parents
|
||||
for name, item in module_names.items():
|
||||
if item == module:
|
||||
key = '{}_{}'.format(module_idx, name)
|
||||
|
||||
info = OrderedDict()
|
||||
info["id"] = id(module)
|
||||
@@ -65,45 +69,55 @@ def summary(model, x, *args, **kwargs):
|
||||
# ignore Sequential and ModuleList
|
||||
if not module._modules:
|
||||
hooks.append(module.register_forward_hook(hook))
|
||||
|
||||
module_names = get_names_dict(model)
|
||||
|
||||
hooks = []
|
||||
summary = OrderedDict()
|
||||
|
||||
model.apply(register_hook)
|
||||
|
||||
with torch.no_grad():
|
||||
model(x) if not (kwargs or args) else model(x, *args, **kwargs)
|
||||
|
||||
for hook in hooks:
|
||||
hook.remove()
|
||||
|
||||
print("-"*100)
|
||||
print("{:<15} {:>20} {:>20} {:>20} {:>20}"
|
||||
.format("Layer", "Kernel Shape", "Output Shape",
|
||||
"# Params (K)", "# Mult-Adds (M)"))
|
||||
|
||||
# Use pandas to align the columns
|
||||
df = pd.DataFrame(summary).T
|
||||
df['Mult-Adds (M)'] = pd.to_numeric(df['macs'], errors='coerce')/1e6
|
||||
df['Params (K)'] = pd.to_numeric(df['params'], errors='coerce')/1e3
|
||||
df = df.rename(columns=dict(
|
||||
ksize='Kernel Shape',
|
||||
out='Output Shape',
|
||||
))
|
||||
df.index.name = 'Layer'
|
||||
df = df[['Kernel Shape', 'Output Shape', 'Params (K)', 'Mult-Adds (M)']]
|
||||
|
||||
print("="*100)
|
||||
print(df.replace(np.nan, '-')
|
||||
print("="*100)
|
||||
print(df.sum())
|
||||
print("="*100)
|
||||
|
||||
total_params, total_macs = 0, 0
|
||||
for layer, info in summary.items():
|
||||
repr_ksize = str(info["ksize"])
|
||||
repr_out = str(info["out"])
|
||||
repr_params = info["params"]
|
||||
repr_macs = info["macs"]
|
||||
return df
|
||||
|
||||
def get_names_dict(model):
|
||||
"""Recursive walk to get names including path."""
|
||||
names = {}
|
||||
|
||||
if isinstance(repr_params, (int, float)):
|
||||
total_params += repr_params
|
||||
repr_params = "{0:,.2f}".format(repr_params/1000)
|
||||
if isinstance(repr_macs, (int, float)):
|
||||
total_macs += repr_macs
|
||||
repr_macs = "{0:,.2f}".format(repr_macs/1000000)
|
||||
|
||||
print("{:<15} {:>20} {:>20} {:>20} {:>20}"
|
||||
.format(layer, repr_ksize, repr_out, repr_params, repr_macs))
|
||||
|
||||
# for RNN, describe inner weights (i.e. w_hh, w_ih)
|
||||
for inner_name, inner_shape in info["inner"].items():
|
||||
print(" {:<13} {:>20}".format(inner_name, str(inner_shape)))
|
||||
|
||||
print("="*100)
|
||||
print("# Params: {0:,.2f}K".format(total_params/1000))
|
||||
print("# Mult-Adds: {0:,.2f}M".format(total_macs/1000000))
|
||||
print("-"*100)
|
||||
def _get_names(module, parent_name=''):
|
||||
for key, module in module.named_children():
|
||||
cls_name = str(module.__class__).split(".")[-1].split("'")[0]
|
||||
num_named_children = len(list(module.named_children()))
|
||||
if num_named_children>0:
|
||||
name = parent_name + '.' + key if parent_name else key
|
||||
else:
|
||||
name = parent_name + '.' + cls_name + '_'+ key if parent_name else key
|
||||
names[name] = module
|
||||
|
||||
if isinstance(module, torch.nn.Module):
|
||||
_get_names(module, parent_name=name)
|
||||
|
||||
_get_names(model)
|
||||
return names
|
||||
|
||||
Reference in New Issue
Block a user