Lyken17 / pytorch-OpCounter

Count the MACs / FLOPs of your PyTorch model.

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

[BUG Report] thop not support AdaptiveAvgPool2d when None in the shape parameter

buptweixin opened this issue · comments

When I set the one dim to None in AdaptiveAvgPool2d to make the dim same as that of the input as described in https://pytorch.org/docs/master/generated/torch.nn.AdaptiveAvgPool2d.html#adaptiveavgpool2d, the thop lib will crash

  File ~.local/lib/python3.8/site-packages/thop/vision/basic_hooks.py", line 92, in count_adap_avgpool
    kernel = torch.DoubleTensor([*(x[0].shape[2:])]) // torch.DoubleTensor(list((m.output_size,))).squeeze()
TypeError: must be real number, not NoneType

I also encountered this problem. I tried to change the code without reporting an error. You can also try it. Just change the 84 lines in site-packages/thop/vision/basic_hooks.py to the following code:

def count_adap_avgpool(m, x, y):
# start-------add by like
   t = [*(x[0].shape[2:])]
   m_out_size = t
   if m.output_size[0] == None:
       m_out_size[1] == 1
   elif m.output_size[1] == None:
       m_out_size[0] == 1
   else:
       m_out_size = m.output_size
   a = torch.Tensor(list((m_out_size ,))).squeeze()
   # end---------add by like
   kernel = torch.Tensor(t)/a # torch.Tensor([*(x[0].shape[2:])]) // torch.Tensor(list((m.output_size,))).squeeze()