diff --git a/python/paddle/hapi/dynamic_flops.py b/python/paddle/hapi/dynamic_flops.py index 382227ea8329778d767a47114496d9365e6def8c..bfbb483ac31ea0537f7f5d2cccdbdd384b3ae444 100644 --- a/python/paddle/hapi/dynamic_flops.py +++ b/python/paddle/hapi/dynamic_flops.py @@ -221,7 +221,8 @@ def dynamic_flops(model, inputs, custom_ops=None, print_detail=False): if m_type in custom_ops: flops_fn = custom_ops[m_type] if m_type not in types_collection: - print("Customize Function has been appied to {}".format(m_type)) + print("Customize Function has been applied to {}".format( + m_type)) elif m_type in register_hooks: flops_fn = register_hooks[m_type] if m_type not in types_collection: @@ -229,7 +230,7 @@ def dynamic_flops(model, inputs, custom_ops=None, print_detail=False): else: if m_type not in types_collection: print( - "Cannot find suitable count function for {}. Treat it as zero Macs.". + "Cannot find suitable count function for {}. Treat it as zero FLOPs.". format(m_type)) if flops_fn is not None: @@ -254,11 +255,10 @@ def dynamic_flops(model, inputs, custom_ops=None, print_detail=False): for m in model.sublayers(): if len(list(m.children())) > 0: continue - total_ops += m.total_ops - total_params += m.total_params - - total_ops = int(total_ops) - total_params = int(total_params) + if set(['total_ops', 'total_params', 'input_shape', + 'output_shape']).issubset(set(list(m._buffers.keys()))): + total_ops += m.total_ops + total_params += m.total_params if training: model.train() @@ -277,7 +277,8 @@ def dynamic_flops(model, inputs, custom_ops=None, print_detail=False): for n, m in model.named_sublayers(): if len(list(m.children())) > 0: continue - if "total_ops" in m._buffers: + if set(['total_ops', 'total_params', 'input_shape', + 'output_shape']).issubset(set(list(m._buffers.keys()))): table.add_row([ m.full_name(), list(m.input_shape.numpy()), list(m.output_shape.numpy()), int(m.total_params), @@ -289,6 +290,6 @@ def dynamic_flops(model, inputs, custom_ops=None, print_detail=False): m._buffers.pop('output_shape') if (print_detail): print(table) - print('Total Flops: {} Total Params: {}'.format(total_ops, - total_params)) - return total_ops + print('Total Flops: {} Total Params: {}'.format( + int(total_ops), int(total_params))) + return int(total_ops)