Skip to content
Permalink
Browse files

add ModelStats, merge PRs for torchstat, add support for adaptive pool

  • Loading branch information
sytelus committed Jan 10, 2020
1 parent fbf9feb commit 44929aaf97f1bb6b162f222256019e726758615f

Large diffs are not rendered by default.

@@ -24,7 +24,7 @@
),
include_package_data=True,
install_requires=[
'matplotlib', 'numpy', 'pyzmq', 'plotly', 'torchstat', 'ipywidgets', 'pydot',
'matplotlib', 'numpy', 'pyzmq', 'plotly', 'ipywidgets', 'pydot',
'nbformat', 'scikit-image', 'nbformat', 'pyyaml', 'scikit-image', 'graphviz' # , 'receptivefield'
]
)
@@ -0,0 +1,7 @@
# Credits

Code in this folder is almost as-is from torchstat repository located at https://github.com/Swall0w/torchstat.

Additional merges are from:
- https://github.com/kenshohara/torchstat
- https://github.com/lyakaap/torchstat
No changes.
@@ -0,0 +1,119 @@
import torch.nn as nn
import torch
import numpy as np
import math

def compute_flops(module, inp, out):
if isinstance(module, nn.Conv2d):
return compute_Conv2d_flops(module, inp, out)
elif isinstance(module, nn.BatchNorm2d):
return compute_BatchNorm2d_flops(module, inp, out)
elif isinstance(module, (nn.AvgPool2d, nn.MaxPool2d)):
return compute_Pool2d_flops(module, inp, out)
elif isinstance(module, (nn.AdaptiveAvgPool2d, nn.AdaptiveMaxPool2d)):
return compute_adaptivepool_flops(module, inp, out)
elif isinstance(module, (nn.ReLU, nn.ReLU6, nn.PReLU, nn.ELU, nn.LeakyReLU)):
return compute_ReLU_flops(module, inp, out)
elif isinstance(module, nn.Upsample):
return compute_Upsample_flops(module, inp, out)
elif isinstance(module, nn.Linear):
return compute_Linear_flops(module, inp, out)
else:
print(f"[Flops]: {type(module).__name__} is not supported!")
return 0
pass


def compute_Conv2d_flops(module, inp, out):
# Can have multiple inputs, getting the first one
assert isinstance(module, nn.Conv2d)
assert len(inp.size()) == 4 and len(inp.size()) == len(out.size())

batch_size = inp.size()[0]
in_c = inp.size()[1]
k_h, k_w = module.kernel_size
out_c, out_h, out_w = out.size()[1:]
groups = module.groups

filters_per_channel = out_c // groups
conv_per_position_flops = k_h * k_w * in_c * filters_per_channel
active_elements_count = batch_size * out_h * out_w

total_conv_flops = conv_per_position_flops * active_elements_count

bias_flops = 0
if module.bias is not None:
bias_flops = out_c * active_elements_count

total_flops = total_conv_flops + bias_flops
return total_flops

def compute_adaptivepool_flops(module, input, output):
# credits: https://github.com/xternalz/SDPoint/blob/master/utils/flops.py
batch_size = input.size(0)
input_planes = input.size(1)
input_height = input.size(2)
input_width = input.size(3)

flops = 0
for i in range(output.size(2)):
y_start = int(math.floor(float(i * input_height) / output.size(2)))
y_end = int(math.ceil(float((i + 1) * input_height) / output.size(2)))
for j in range(output.size(3)):
x_start = int(math.floor(float(j * input_width) / output.size(3)))
x_end = int(math.ceil(float((j + 1) * input_width) / output.size(3)))

flops += batch_size * input_planes * (y_end-y_start+1) * (x_end-x_start+1)
return flops

def compute_BatchNorm2d_flops(module, inp, out):
assert isinstance(module, nn.BatchNorm2d)
assert len(inp.size()) == 4 and len(inp.size()) == len(out.size())
in_c, in_h, in_w = inp.size()[1:]
batch_flops = np.prod(inp.shape)
if module.affine:
batch_flops *= 2
return batch_flops


def compute_ReLU_flops(module, inp, out):
assert isinstance(module, (nn.ReLU, nn.ReLU6, nn.PReLU, nn.ELU, nn.LeakyReLU))
batch_size = inp.size()[0]
active_elements_count = batch_size

for s in inp.size()[1:]:
active_elements_count *= s

return active_elements_count


def compute_Pool2d_flops(module, input, out):
batch_size = input.size(0)
input_planes = input.size(1)
input_height = input.size(2)
input_width = input.size(3)
kernel_size = ('int' in str(type(module.kernel_size))) and [module.kernel_size, module.kernel_size] or module.kernel_size
kernel_ops = kernel_size[0] * kernel_size[1]
stride = ('int' in str(type(module.stride))) and [module.stride, module.stride] or module.stride
padding = ('int' in str(type(module.padding))) and [module.padding, module.padding] or module.padding

output_width = math.floor((input_width + 2 * padding[0] - kernel_size[0]) / float(stride[0]) + 1)
output_height = math.floor((input_height + 2 * padding[1] - kernel_size[1]) / float(stride[0]) + 1)
return batch_size * input_planes * output_width * output_height * kernel_ops


def compute_Linear_flops(module, inp, out):
assert isinstance(module, nn.Linear)
assert len(inp.size()) == 2 and len(out.size()) == 2
batch_size = inp.size()[0]
return batch_size * inp.size()[1] * out.size()[1]

def compute_Upsample_flops(module, inp, out):
assert isinstance(module, nn.Upsample)
output_size = out[0]
batch_size = inp.size()[0]
output_elements_count = batch_size
for s in output_size.shape[1:]:
output_elements_count *= s

return output_elements_count
@@ -0,0 +1,161 @@
"""
compute Multiply-Adds(MAdd) of each leaf module
"""

import torch.nn as nn


def compute_Conv2d_madd(module, inp, out):
assert isinstance(module, nn.Conv2d)
assert len(inp.size()) == 4 and len(inp.size()) == len(out.size())

in_c = inp.size()[1]
k_h, k_w = module.kernel_size
out_c, out_h, out_w = out.size()[1:]
groups = module.groups

# ops per output element
kernel_mul = k_h * k_w * (in_c // groups)
kernel_add = kernel_mul - 1 + (0 if module.bias is None else 1)

kernel_mul_group = kernel_mul * out_h * out_w * (out_c // groups)
kernel_add_group = kernel_add * out_h * out_w * (out_c // groups)

total_mul = kernel_mul_group * groups
total_add = kernel_add_group * groups

return total_mul + total_add


def compute_ConvTranspose2d_madd(module, inp, out):
assert isinstance(module, nn.ConvTranspose2d)
assert len(inp.size()) == 4 and len(inp.size()) == len(out.size())

in_c, in_h, in_w = inp.size()[1:]
k_h, k_w = module.kernel_size
out_c, out_h, out_w = out.size()[1:]
groups = module.groups

kernel_mul = k_h * k_w * (in_c // groups)
kernel_add = kernel_mul - 1 + (0 if module.bias is None else 1)

kernel_mul_group = kernel_mul * in_h * in_w * (out_c // groups)
kernel_add_group = kernel_add * in_h * in_w * (out_c // groups)

total_mul = kernel_mul_group * groups
total_add = kernel_add_group * groups

return total_mul + total_add


def compute_BatchNorm2d_madd(module, inp, out):
assert isinstance(module, nn.BatchNorm2d)
assert len(inp.size()) == 4 and len(inp.size()) == len(out.size())

in_c, in_h, in_w = inp.size()[1:]

# 1. sub mean
# 2. div standard deviation
# 3. mul alpha
# 4. add beta
return 4 * in_c * in_h * in_w


def compute_MaxPool2d_madd(module, inp, out):
assert isinstance(module, nn.MaxPool2d)
assert len(inp.size()) == 4 and len(inp.size()) == len(out.size())

if isinstance(module.kernel_size, (tuple, list)):
k_h, k_w = module.kernel_size
else:
k_h, k_w = module.kernel_size, module.kernel_size
out_c, out_h, out_w = out.size()[1:]

return (k_h * k_w - 1) * out_h * out_w * out_c


def compute_AvgPool2d_madd(module, inp, out):
assert isinstance(module, nn.AvgPool2d)
assert len(inp.size()) == 4 and len(inp.size()) == len(out.size())

if isinstance(module.kernel_size, (tuple, list)):
k_h, k_w = module.kernel_size
else:
k_h, k_w = module.kernel_size, module.kernel_size
out_c, out_h, out_w = out.size()[1:]

kernel_add = k_h * k_w - 1
kernel_avg = 1

return (kernel_add + kernel_avg) * (out_h * out_w) * out_c


def compute_ReLU_madd(module, inp, out):
assert isinstance(module, (nn.ReLU, nn.ReLU6))

count = 1
for i in inp.size()[1:]:
count *= i
return count


def compute_Softmax_madd(module, inp, out):
assert isinstance(module, nn.Softmax)
assert len(inp.size()) > 1

count = 1
for s in inp.size()[1:]:
count *= s
exp = count
add = count - 1
div = count
return exp + add + div


def compute_Linear_madd(module, inp, out):
assert isinstance(module, nn.Linear)
assert len(inp.size()) == 2 and len(out.size()) == 2

num_in_features = inp.size()[1]
num_out_features = out.size()[1]

mul = num_in_features
add = num_in_features - 1
return num_out_features * (mul + add)


def compute_Bilinear_madd(module, inp1, inp2, out):
assert isinstance(module, nn.Bilinear)
assert len(inp1.size()) == 2 and len(inp2.size()) == 2 and len(out.size()) == 2

num_in_features_1 = inp1.size()[1]
num_in_features_2 = inp2.size()[1]
num_out_features = out.size()[1]

mul = num_in_features_1 * num_in_features_2 + num_in_features_2
add = num_in_features_1 * num_in_features_2 + num_in_features_2 - 1
return num_out_features * (mul + add)


def compute_madd(module, inp, out):
if isinstance(module, nn.Conv2d):
return compute_Conv2d_madd(module, inp, out)
elif isinstance(module, nn.ConvTranspose2d):
return compute_ConvTranspose2d_madd(module, inp, out)
elif isinstance(module, nn.BatchNorm2d):
return compute_BatchNorm2d_madd(module, inp, out)
elif isinstance(module, nn.MaxPool2d):
return compute_MaxPool2d_madd(module, inp, out)
elif isinstance(module, nn.AvgPool2d):
return compute_AvgPool2d_madd(module, inp, out)
elif isinstance(module, (nn.ReLU, nn.ReLU6)):
return compute_ReLU_madd(module, inp, out)
elif isinstance(module, nn.Softmax):
return compute_Softmax_madd(module, inp, out)
elif isinstance(module, nn.Linear):
return compute_Linear_madd(module, inp, out)
elif isinstance(module, nn.Bilinear):
return compute_Bilinear_madd(module, inp[0], inp[1], out)
else:
print(f"[MAdd]: {type(module).__name__} is not supported!")
return 0
@@ -0,0 +1,92 @@
import torch.nn as nn
import torch
import numpy as np


def compute_memory(module, inp, out):
if isinstance(module, (nn.ReLU, nn.ReLU6, nn.ELU, nn.LeakyReLU)):
return compute_ReLU_memory(module, inp, out)
elif isinstance(module, nn.PReLU):
return compute_PReLU_memory(module, inp, out)
elif isinstance(module, nn.Conv2d):
return compute_Conv2d_memory(module, inp, out)
elif isinstance(module, nn.BatchNorm2d):
return compute_BatchNorm2d_memory(module, inp, out)
elif isinstance(module, nn.Linear):
return compute_Linear_memory(module, inp, out)
elif isinstance(module, (nn.AvgPool2d, nn.MaxPool2d)):
return compute_Pool2d_memory(module, inp, out)
else:
print(f"[Memory]: {type(module).__name__} is not supported!")
return 0, 0
pass


def num_params(module):
return sum(p.numel() for p in module.parameters() if p.requires_grad)


def compute_ReLU_memory(module, inp, out):
assert isinstance(module, (nn.ReLU, nn.ReLU6, nn.ELU, nn.LeakyReLU))

mread = inp.numel()
mwrite = out.numel()

return mread, mwrite


def compute_PReLU_memory(module, inp, out):
assert isinstance(module, nn.PReLU)

batch_size = inp.size()[0]
mread = batch_size * (inp[0].numel() + num_params(module))
mwrite = out.numel()

return mread, mwrite


def compute_Conv2d_memory(module, inp, out):
# Can have multiple inputs, getting the first one
assert isinstance(module, nn.Conv2d)
assert len(inp.size()) == 4 and len(inp.size()) == len(out.size())

batch_size = inp.size()[0]

# This includes weights with bias if the module contains it.
mread = batch_size * (inp[0].numel() + num_params(module))
mwrite = out.numel()
return mread, mwrite


def compute_BatchNorm2d_memory(module, inp, out):
assert isinstance(module, nn.BatchNorm2d)
assert len(inp.size()) == 4 and len(inp.size()) == len(out.size())

batch_size, in_c, in_h, in_w = inp.size()
mread = batch_size * (inp[0].numel() + 2 * in_c)
mwrite = out.numel()

return mread, mwrite


def compute_Linear_memory(module, inp, out):
assert isinstance(module, nn.Linear)
assert len(inp.size()) == 2 and len(out.size()) == 2

batch_size = inp.size()[0]

# This includes weights with bias if the module contains it.
mread = batch_size * (inp[0].numel() + num_params(module))
mwrite = out.numel()

return mread, mwrite


def compute_Pool2d_memory(module, inp, out):
assert isinstance(module, (nn.MaxPool2d, nn.AvgPool2d))
assert len(inp.size()) == 4 and len(inp.size()) == len(out.size())

mread = inp.numel()
mwrite = out.numel()

return mread, mwrite

0 comments on commit 44929aa

Please sign in to comment.
You can’t perform that action at this time.