From 68fb998dc3d27d77de7b3cdcc6fab9c2c5c5849e Mon Sep 17 00:00:00 2001 From: ceci3 <592712189@qq.com> Date: Tue, 7 Jul 2020 06:38:24 +0000 Subject: [PATCH] add macs --- paddleslim/analysis/flops.py | 99 ++++++++++++++++++++++++++---------- 1 file changed, 73 insertions(+), 26 deletions(-) diff --git a/paddleslim/analysis/flops.py b/paddleslim/analysis/flops.py index 59154610..ab58ba48 100644 --- a/paddleslim/analysis/flops.py +++ b/paddleslim/analysis/flops.py @@ -16,13 +16,12 @@ import numpy as np import paddle.fluid.dygraph.jit as jit from ..core import GraphWrapper -__all__ = ["dygraph_flops", "flops"] +__all__ = ["dygraph_flops", "flops", "dygrpah_macs", "macs"] def dygraph_flops(model, inputs, only_conv=True, - only_multiply=False, detail=False): """ Get FLOPs of dygraph model. @@ -31,26 +30,21 @@ def dygraph_flops(model, inputs: The inputs of the model, used to calculate FLOPs. only_conv(bool): Just return number of mul-adds in convolution and FC layer if `only_conv` is true. Default: True. - only_multiply(bool): If `only_multiply` is true, just return number of muliply in the model, the - multiply in such as conv, conv_transpose, norm and mul operators will be count. - Default: False. detail(bool): Whether to return detail of each convolution layer. Default: False. """ _, program, _, _, _ = jit._trace(model, inputs) graph = GraphWrapper(program) - return _graph_flops( - graph, only_conv=only_conv, only_multiply=only_multiply, detail=detail) + return _graph_cals( + graph, only_conv=only_conv, count='FLOPs', detail=detail) -def flops(program, only_conv=True, only_multiply=False, detail=False): +def flops(program, only_conv=True, detail=False): """ Get FLOPS of target graph. Args: program(Program): The program used to calculate FLOPS. only_conv(bool): Just return number of mul-adds in convolution and FC layer if `only_conv` is true. default: True. - only_multiply(bool): Just return number of muliply in the model if `only_multiply` is true. - Default: False. detail(bool): Whether to return detail of each convolution layer. Default: False. Return: @@ -59,12 +53,50 @@ def flops(program, only_conv=True, only_multiply=False, detail=False): details(dict): The key is the parameter name of convlution layer and the value is the FLOPs of each convolution layer. """ graph = GraphWrapper(program) + return _graph_cals( + graph, only_conv=only_conv, count='FLOPs', detail=detail) + + +def dygraph_macs(model, + inputs, + only_conv=False, + detail=False): + """ + Get FLOPs of dygraph model. + Args: + model: The dygraph model to calculate FLOPs. + inputs: The inputs of the model, used to calculate FLOPs. + only_conv(bool): Just return number of mul-adds in convolution and FC layer if `only_conv` is true. + Default: True. + detail(bool): Whether to return detail of each convolution layer. Default: False. + """ + _, program, _, _, _ = jit._trace(model, inputs) + graph = GraphWrapper(program) return _graph_flops( - graph, only_conv=only_conv, only_multiply=only_multiply, detail=detail) + graph, only_conv=only_conv, count='MACs', detail=detail) -def _graph_flops(graph, only_conv=True, only_multiply=False, detail=False): +def macs(program, only_conv=True, detail=False): + """ + Get FLOPS of target graph. + Args: + program(Program): The program used to calculate FLOPS. + only_conv(bool): Just return number of mul-adds in convolution and FC layer if `only_conv` is true. + default: True. + detail(bool): Whether to return detail of each convolution layer. Default: False. + + Return: + If `detail` is true, then return a tuple in format `(FLOPs, details)`, otherwise it will just return `FlOPs` + FLOPs(int): The FLOPs of target network. + details(dict): The key is the parameter name of convlution layer and the value is the FLOPs of each convolution layer. + """ + graph = GraphWrapper(program) + return _graph_cals( + graph, only_conv=only_conv, count='MACs', detail=detail) + +def _graph_cals(graph, only_conv=True, count, detail=False): assert isinstance(graph, GraphWrapper) + assert count.lower() in ['flops', 'macs'], "count {} not support now".format(count) flops = 0 params2flops = {} for op in graph.ops(): @@ -77,32 +109,43 @@ def _graph_flops(graph, only_conv=True, only_multiply=False, detail=False): c_out, c_in, k_h, k_w = filter_shape _, _, h_out, w_out = output_shape # c_in is the channel number of filter. It is (input_channel // groups). - kernel_ops = k_h * k_w * float(c_in) - ### after dygraph model to static program, conv op donnot have Bias attrs - op_flops = h_out * w_out * c_out * kernel_ops + if count == 'MACs': + out_pixel = k_h * k_w * float(c_in) + else: + ### count == 'FLOPs' + ### add bias count in elementwise_add op + out_pixel = 2 * k_h * k_w * float(c_in) - 1 + + op_flops = h_out * w_out * c_out * out_pixel flops += op_flops params2flops[op.inputs("Filter")[0].name()] = op_flops - elif op.type() == 'pool2d' and not only_conv and not only_multiply: - output_shape = op.outputs("Out")[0].shape() - _, c_out, h_out, w_out = output_shape + elif op.type() == 'pool2d' and not only_conv: + if count == 'MACs': + op_shape = op.outputs("Out")[0].shape() + else: + op_shape = op.inputs("X")[0].shape() + _, c_out, h_out, w_out = op_shape k_size = op.attr("ksize") op_flops = h_out * w_out * c_out * (k_size[0]**2) flops += op_flops - elif op.type() == 'mul': + elif op.type() == 'mul' and not only_conv: x_shape = list(op.inputs("X")[0].shape()) y_shape = op.inputs("Y")[0].shape() if x_shape[0] == -1: x_shape[0] = 1 - flops += x_shape[0] * x_shape[1] * y_shape[1] - op_flops = x_shape[0] * x_shape[1] * y_shape[1] + if count == 'MACs': + op_flops = x_shape[0] * x_shape[1] * y_shape[1] + else: + op_flops = (2 * x_shape[0] * x_shape[1] - 1) * y_shape[1] + flops += op_flops params2flops[op.inputs("Y")[0].name()] = op_flops elif op.type() in ['relu', 'sigmoid', 'relu6' - ] and not only_conv and not only_multiply: + ] and not only_conv: input_shape = list(op.inputs("X")[0].shape()) if input_shape[0] == -1: input_shape[0] = 1 @@ -110,17 +153,21 @@ def _graph_flops(graph, only_conv=True, only_multiply=False, detail=False): flops += op_flops elif op.type() in ['batch_norm', 'instance_norm', 'layer_norm' - ] and not only_conv and only_multiply: + ] and not only_conv: input_shape = list(op.inputs("X")[0].shape()) if input_shape[0] == -1: input_shape[0] = 1 + ### (x - mean) * sqrt(var) op_flops = np.product(input_shape) + if count == 'FLOPs': + ### NOTE: if scale and bias can be none (Need to fix in bn op), it need to add more condition to determine if need to multiply 2 + op_flops *= 2 flops += op_flops - elif op.type() in ['elementwise_add'] and not only_multiply: + elif op.type() in ['elementwise_add'] and count == 'FLOPs': input_shape = list(op.inputs("X")[0].shape()) - ### if inputs Y is parameter that means add bias after conv - if op.inputs("Y")[0].is_parameter() or not only_conv: + ### if inputs Y is parameter that means add bias after conv or norm + if op.inputs("Y")[0].is_parameter(): if input_shape[0] == -1: input_shape[0] = 1 op_flops = np.product(input_shape) -- GitLab