You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
195 lines
7.1 KiB
195 lines
7.1 KiB
# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
'''
|
|
Example:
|
|
>>from paddle.fluid.contrib.model_stat import summary
|
|
>>main_program = ...
|
|
>>summary(main_program)
|
|
+-----+------------+----------------+----------------+---------+------------+
|
|
| No. | TYPE | INPUT | OUTPUT | PARAMs | FLOPs |
|
|
+-----+------------+----------------+----------------+---------+------------+
|
|
| 0 | conv2d | (3, 200, 200) | (64, 100, 100) | 9408 | 188160000 |
|
|
| 1 | batch_norm | (64, 100, 100) | (64, 100, 100) | 256 | 640000 |
|
|
| 2 | relu | (64, 100, 100) | (64, 100, 100) | 0 | 640000 |
|
|
| 3 | pool2d | (64, 100, 100) | (64, 50, 50) | 0 | 1440000 |
|
|
...
|
|
| 176 | conv2d | (512, 7, 7) | (512, 7, 7) | 2359296 | 231211008 |
|
|
| 177 | relu | (512, 7, 7) | (512, 7, 7) | 0 | 25088 |
|
|
| 178 | conv2d | (512, 7, 7) | (2048, 7, 7) | 1048576 | 102760448 |
|
|
| 179 | relu | (2048, 7, 7) | (2048, 7, 7) | 0 | 100352 |
|
|
| 180 | pool2d | (2048, 7, 7) | (2048, 1, 1) | 0 | 100352 |
|
|
+-----+------------+----------------+----------------+---------+------------+
|
|
Total PARAMs: 48017344(0.0480G)
|
|
Total FLOPs: 11692747751(11.69G)
|
|
'''
|
|
from collections import OrderedDict
|
|
from prettytable import PrettyTable
|
|
|
|
|
|
def summary(main_prog):
|
|
'''
|
|
It can summary model's PARAMS, FLOPs until now.
|
|
It support common operator like conv, fc, pool, relu, sigmoid, bn etc.
|
|
Args:
|
|
main_prog: main program
|
|
Returns:
|
|
print summary on terminal
|
|
'''
|
|
collected_ops_list = []
|
|
for one_b in main_prog.blocks:
|
|
block_vars = one_b.vars
|
|
for one_op in one_b.ops:
|
|
op_info = OrderedDict()
|
|
spf_res = _summary_model(block_vars, one_op)
|
|
if spf_res is None:
|
|
continue
|
|
# TODO: get the operator name
|
|
op_info['type'] = one_op.type
|
|
op_info['input_shape'] = spf_res[0][1:]
|
|
op_info['out_shape'] = spf_res[1][1:]
|
|
op_info['PARAMs'] = spf_res[2]
|
|
op_info['FLOPs'] = spf_res[3]
|
|
collected_ops_list.append(op_info)
|
|
|
|
summary_table, total = _format_summary(collected_ops_list)
|
|
_print_summary(summary_table, total)
|
|
|
|
|
|
def _summary_model(block_vars, one_op):
|
|
'''
|
|
Compute operator's params and flops.
|
|
Args:
|
|
block_vars: all vars of one block
|
|
one_op: one operator to count
|
|
Returns:
|
|
in_data_shape: one operator's input data shape
|
|
out_data_shape: one operator's output data shape
|
|
params: one operator's PARAMs
|
|
flops: : one operator's FLOPs
|
|
'''
|
|
if one_op.type in ['conv2d', 'depthwise_conv2d']:
|
|
k_arg_shape = block_vars[one_op.input("Filter")[0]].shape
|
|
in_data_shape = block_vars[one_op.input("Input")[0]].shape
|
|
out_data_shape = block_vars[one_op.output("Output")[0]].shape
|
|
c_out, c_in, k_h, k_w = k_arg_shape
|
|
_, c_out_, h_out, w_out = out_data_shape
|
|
assert c_out == c_out_, 'shape error!'
|
|
k_groups = one_op.attr("groups")
|
|
kernel_ops = k_h * k_w * (c_in / k_groups)
|
|
bias_ops = 0 if one_op.input("Bias") == [] else 1
|
|
params = c_out * (kernel_ops + bias_ops)
|
|
flops = h_out * w_out * c_out * (kernel_ops + bias_ops)
|
|
# base nvidia paper, include mul and add
|
|
flops = 2 * flops
|
|
|
|
elif one_op.type == 'pool2d':
|
|
in_data_shape = block_vars[one_op.input("X")[0]].shape
|
|
out_data_shape = block_vars[one_op.output("Out")[0]].shape
|
|
_, c_out, h_out, w_out = out_data_shape
|
|
k_size = one_op.attr("ksize")
|
|
params = 0
|
|
flops = h_out * w_out * c_out * (k_size[0] * k_size[1])
|
|
|
|
elif one_op.type == 'mul':
|
|
k_arg_shape = block_vars[one_op.input("Y")[0]].shape
|
|
in_data_shape = block_vars[one_op.input("X")[0]].shape
|
|
out_data_shape = block_vars[one_op.output("Out")[0]].shape
|
|
# TODO: fc has mul ops
|
|
# add attr to mul op, tell us whether it belongs to 'fc'
|
|
# this's not the best way
|
|
if 'fc' not in one_op.output("Out")[0]:
|
|
return None
|
|
k_in, k_out = k_arg_shape
|
|
# bias in sum op
|
|
params = k_in * k_out + 1
|
|
flops = k_in * k_out
|
|
|
|
elif one_op.type in ['sigmoid', 'tanh', 'relu', 'leaky_relu', 'prelu']:
|
|
in_data_shape = block_vars[one_op.input("X")[0]].shape
|
|
out_data_shape = block_vars[one_op.output("Out")[0]].shape
|
|
params = 0
|
|
if one_op.type == 'prelu':
|
|
params = 1
|
|
flops = 1
|
|
for one_dim in in_data_shape:
|
|
flops *= one_dim
|
|
|
|
elif one_op.type == 'batch_norm':
|
|
in_data_shape = block_vars[one_op.input("X")[0]].shape
|
|
out_data_shape = block_vars[one_op.output("Y")[0]].shape
|
|
_, c_in, h_out, w_out = in_data_shape
|
|
# gamma, beta
|
|
params = c_in * 2
|
|
# compute mean and std
|
|
flops = h_out * w_out * c_in * 2
|
|
|
|
else:
|
|
return None
|
|
|
|
return in_data_shape, out_data_shape, params, flops
|
|
|
|
|
|
def _format_summary(collected_ops_list):
|
|
'''
|
|
Format summary report.
|
|
Args:
|
|
collected_ops_list: the collected operator with summary
|
|
Returns:
|
|
summary_table: summary report format
|
|
total: sum param and flops
|
|
'''
|
|
summary_table = PrettyTable(
|
|
["No.", "TYPE", "INPUT", "OUTPUT", "PARAMs", "FLOPs"])
|
|
summary_table.align = 'r'
|
|
|
|
total = {}
|
|
total_params = []
|
|
total_flops = []
|
|
for i, one_op in enumerate(collected_ops_list):
|
|
# notice the order
|
|
table_row = [
|
|
i,
|
|
one_op['type'],
|
|
one_op['input_shape'],
|
|
one_op['out_shape'],
|
|
int(one_op['PARAMs']),
|
|
int(one_op['FLOPs']),
|
|
]
|
|
summary_table.add_row(table_row)
|
|
total_params.append(int(one_op['PARAMs']))
|
|
total_flops.append(int(one_op['FLOPs']))
|
|
|
|
total['params'] = total_params
|
|
total['flops'] = total_flops
|
|
|
|
return summary_table, total
|
|
|
|
|
|
def _print_summary(summary_table, total):
|
|
'''
|
|
Print all the summary on terminal.
|
|
Args:
|
|
summary_table: summary report format
|
|
total: sum param and flops
|
|
'''
|
|
parmas = total['params']
|
|
flops = total['flops']
|
|
print(summary_table)
|
|
print('Total PARAMs: {}({:.4f}M)'.format(
|
|
sum(parmas), sum(parmas) / (10**6)))
|
|
print('Total FLOPs: {}({:.2f}G)'.format(sum(flops), sum(flops) / 10**9))
|
|
print(
|
|
"Notice: \n now supported ops include [Conv, DepthwiseConv, FC(mul), BatchNorm, Pool, Activation(sigmoid, tanh, relu, leaky_relu, prelu)]"
|
|
)
|