set net.forward to original forward function in flops

Author: 0x45f

python/paddle/hapi/dynamic_flops.py

@@ -17,6 +17,7 @@
 import paddle.nn as nn
 import numpy as np
 from .static_flops import static_flops, Table
+from paddle.fluid.dygraph.dygraph_to_static.program_translator import unwrap_decorators
 
 __all__ = []
 
@@ -100,6 +101,10 @@ def count_leaky_relu(m, x, y):
  #Total Flops: 347560 Total Params: 61610
  """
  if isinstance(net, nn.Layer):
+ # If net is a dy2stat model, net.forward is StaticFunction instance,
+ # we set net.forward to original forward function.
+ _, net.forward = unwrap_decorators(net.forward)
+
  inputs = paddle.randn(input_size)
  return dynamic_flops(
  net,