[BUG Report] thop not support AdaptiveAvgPool2d when None in the shape parameter
buptweixin opened this issue · comments
Arno Wei commented
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
like121 commented
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()