forked from fxmeng/Pruning-Filter-in-Filter
-
Notifications
You must be signed in to change notification settings - Fork 0
/
flops.py
92 lines (68 loc) · 3.24 KB
/
flops.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
# Code from https://github.com/simochen/model-tools.
import numpy as np
import os
from models.stripe import *
import torch
import torchvision
import torch.nn as nn
def print_model_param_nums(model=None):
total = sum([param.nelement() for param in model.parameters()])
print(' + Number of params: %.2fM' % (total / 1e6))
def count_model_param_flops(model=None, input_res=32, multiply_adds=True):
list_strip = []
def filter_strip_hook(self, input, output):
output_channels, output_height, output_width = output[0].size()
flops = self.kernel_size[0] * self.kernel_size[1] * self.in_channels * (2 if multiply_adds else 1) * self.out_channels * output_height * output_width * input[0].size()[0]
list_strip.append(flops)
list_linear = []
def linear_hook(self, input, output):
batch_size = input[0].size(0) if input[0].dim() == 2 else 1
weight_ops = self.weight.nelement() * (2 if multiply_adds else 1)
bias_ops = self.bias.nelement()
flops = batch_size * (weight_ops + bias_ops)
list_linear.append(flops)
list_bn = []
def bn_hook(self, input, output):
list_bn.append(input[0].nelement() * 2)
list_relu = []
def relu_hook(self, input, output):
list_relu.append(input[0].nelement())
list_pooling = []
def pooling_hook(self, input, output):
batch_size, input_channels, input_height, input_width = input[0].size()
output_channels, output_height, output_width = output[0].size()
kernel_ops = self.kernel_size * self.kernel_size
bias_ops = 0
params = 0
flops = (kernel_ops + bias_ops) * output_channels * output_height * output_width * batch_size
list_pooling.append(flops)
list_upsample = []
# For bilinear upsample
def upsample_hook(self, input, output):
batch_size, input_channels, input_height, input_width = input[0].size()
output_channels, output_height, output_width = output[0].size()
flops = output_height * output_width * output_channels * batch_size * 12
list_upsample.append(flops)
def foo(net):
childrens = list(net.children())
if not childrens:
if isinstance(net, FilterStripe):
net.register_forward_hook(filter_strip_hook)
if isinstance(net, Linear):
net.register_forward_hook(linear_hook)
if isinstance(net, BatchNorm):
net.register_forward_hook(bn_hook)
if isinstance(net, torch.nn.ReLU):
net.register_forward_hook(relu_hook)
if isinstance(net, torch.nn.MaxPool2d) or isinstance(net, torch.nn.AvgPool2d):
net.register_forward_hook(pooling_hook)
if isinstance(net, torch.nn.Upsample):
net.register_forward_hook(upsample_hook)
return
for c in childrens:
foo(c)
foo(model)
input = torch.rand(1, 3, input_res, input_res).cuda()
out = model(input)
total_flops = (sum(list_strip) + sum(list_linear) + sum(list_bn) + sum(list_relu) + sum(list_pooling) + sum(list_upsample))
print(' + Number of FLOPs: %.2fM' % (total_flops / 1e6))