Quantizing MaxPool3D

AfonsoSalgadoSousa: <p>Hello. I am trying to employ post-training quantization on a SqueezeNet 3D and am getting the following error:</p>
<blockquote>
<p>RuntimeError: Could not run ‘aten::max_pool3d_with_indices’ with arguments from the ‘QuantizedCPU’ backend. ‘aten::max_pool3d_with_indices’ is only available for these backends: [CPU, CUDA, Named, Autograd, Profiler, Tracer].</p>
</blockquote>
<p>I checked the list of supported operators (<a href="https://pytorch.org/docs/stable/quantization-support.html" class="inline-onebox" rel="noopener nofollow ugc">Quantization API Reference — PyTorch 2.1 documentation</a>) and it seems that MaxPool3D is not supported. How can I circumvent this issue?</p>

Vasiliy_Kuznetsov: <p>if an operator does not have a quantized kernel, you can run it in fp32.  For example,</p>
<pre><code class="lang-auto">class M(torch.nn.Module):
  def __init__(...):
    ...
    self.dequant_1 = torch.quantization.DeQuantStub()
    self.quant_1 = torch.quantization.QuantStub()
    ...

  def forward(self, ...):
    ...
    # (computation in int8)
    ...
    # convert from int8 to fp32
    x_1 = self.dequant_1(x_0)
    # run in fp32
    x_2 = self.maxpool(x_1)
    # convert back to int8, if needed
    x_3 = self.quant_q(x_2)
    ...
    
  
</code></pre>

AfonsoSalgadoSousa: <p>Thank you very much for your reply. That actually answers my question, but hopefully, I can take the opportunity to ask you another one. I am using a pretrained SqueezeNet 3D and in order to be able to load saved models I need to keep the nn.Sequential modules (namely self.features and self.classifier). I tried doing something like this:</p>
<pre><code class="lang-auto">        if self.quantize:
            # run Pooling operations and dropout in FP32
            x = self.quant(x)
            # x = self.features(x)
            for m in self.features:
                if isinstance(m, nn.MaxPool3d):
                    x = self.dequant(x)
                    x = m(x)
                    x = self.quant(x)
                else:
                    x = m(x)
            
            for c in self.classifier:
                if isinstance(c, nn.AvgPool3d) or isinstance(c, nn.Dropout):
                    x = self.dequant(x)
                    x = c(x)
                    x = self.quant(x)
                else:
                    x = c(x)
            x = x.view(x.size(0), -1)
            x = self.dequant(x)
            return x
        else:
            x = self.features(x)
            x = self.classifier(x)
            return x.view(x.size(0), -1)
</code></pre>
<p>This gives the following error:</p>
<blockquote>
<p>RuntimeError:<br>
Cannot re-assign ‘c’ because it has type value of type ‘<strong>torch</strong>.torch.nn.modules.pooling.AvgPool3d’ and c is not a first-class value.  Only reassignments to first-class values are allowed:<br>
File “/home/ctm/afonso/easyride/acceleration/src/models/squeezenet.py”, line 152</p>
<pre><code>        for c in self.classifier:
            if isinstance(c, nn.AvgPool3d) or isinstance(c, nn.Dropout):
               ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ &lt;--- HERE
                x = self.dequant(x)
                x = c(x)
</code></pre>
</blockquote>
<p>Although not understanding why, I actually managed to fix the problem using the following code extract instead of the classifier loop:</p>
<pre><code class="lang-auto">            x = self.dequant(x)
            x = self.classifier[0](x)
            x = self.quant(x)
            x = self.classifier[1](x)
            x = self.classifier[2](x)
            x = self.dequant(x)
            x = self.classifier[-1](x)
</code></pre>
<p>I still do not fully understand how the TorchScript inspector works. I would appreciate if you could give me some insights on this behaviour.</p>