TypeError: avg_pool2d(): argument 'kernel_size' must be tuple of ints, not Proxy

Hello everyone, hope you are having a great time.
I have been playing with quantization recently, and during my last experiments, I tried to change something in my forward pass likethis :

 def forward(self, x):
        out1 = self.p1(x)
        out2 = self.p2(out1)
        out3 = self.p3(out2)

       out1 = out1.mean(3).mean(2)
       out2 = out2.mean(3).mean(2)
       out3 = out3.mean(3).mean(2)
       
       out5 = torch.cat((out1, out2, out3), 1)
       output = self.classifier(out5)
       return output

to

 def forward(self, x):
        out1 = self.p1(x)
        out2 = self.p2(out1)
        out3 = self.p3(out2)

        out1 = F.avg_pool2d(out1, kernel_size=out1.size()[2:]).view(out1.size(0), -1)
        out2 = F.avg_pool2d(out2, kernel_size=out2.size()[2:]).view(out2.size(0), -1)
        out3 = F.avg_pool2d(out3, kernel_size=out3.size()[2:]).view(out3.size(0), -1)

       out5 = torch.cat((out1, out2, out3), 1)
       output = self.classifier(out5)
       return output

I did this so I can easily experiment with qconfig, for example, prevent these operations from participating in the quantization process. something like this :

    qconfig = get_default_qconfig('fbgemm')
    qconfig_dict = {'': qconfig,
                    "object_type": [
                                    (torch.nn.Linear, None),
                                    (torch.nn.functional.avg_pool2d, None)
                                   ]
                    }

But upon trying to start the whole quantization process after this change I get the following error :

    model = torch._fx.symbolic_trace(model)
  File "C:\Users\User\Anaconda3\Lib\site-packages\torch\_fx\symbolic_trace.py", line 168, in symbolic_trace
    return Tracer().trace(root)
  File "C:\Users\User\Anaconda3\Lib\site-packages\torch\_fx\symbolic_trace.py", line 152, in trace
    self.graph.output(self.create_arg(fn(*args)))
  File "d:\Codes\fac_ver\python\Side_Projects\Facial-Landmark-PyTorch\models\slim.py", line 197, in forward
    out1 = F.avg_pool2d(out1, kernel_size=out1.size()[2:]).view(out1.size(0), -1)
TypeError: avg_pool2d(): argument 'kernel_size' must be tuple of ints, not Proxy

and it seems out.size(2:] is Proxy(getitem)
How should I get around this? I know I’m converting the model into symbolic_trace, but if I don’t, then I cant call the quantize_static_fx which requires this.
How should I go about this?

pn: I know I can also use torch.mean just fine, but I’d like to know how to get around that Proxything in general.

Thanks a lot in advance

The fix would be to specify kernel size directly instead of dynamically inferring it from the input tensor. At the moment, this is a restriction of symbolic tracing.

Note: there will be more docs posted on this when FX graph mode quantization is officially released, ideally soon.

1 Like

still not working in pytorch 1.11 and 1.12.
When can we expect this FX limitation to be resolved?

probably not, we are moving away from fx, please look out for news in the end of year regarding PyTorch 2.0, we will be moving quantization stack to PyTorch 2.0 as well

1 Like