Static quatization example not quite working as expected

ResidentMario: <p>I have a <a href="https://gist.github.com/ResidentMario/a29e33a25448f27ee3965b7c0edd4a37" rel="noopener nofollow ugc">quantization script (link to full code here)</a>. Quoting from the code, here are the relevant model layers:</p>
<pre><code class="lang-auto">class UNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.quant_1 = torch.quantization.QuantStub()
        self.conv_1_1 = nn.Conv2d(3, 64, 3)
        torch.nn.init.kaiming_normal_(self.conv_1_1.weight)
        self.relu_1_2 = nn.ReLU()
        self.norm_1_3 = nn.BatchNorm2d(64)
        self.dequant_1 = torch.quantization.DeQuantStub()
        self.conv_1_4 = nn.Conv2d(64, 64, 3)
        torch.nn.init.kaiming_normal_(self.conv_1_4.weight)
        # continued...

    def forward(self, x):
        x = self.quant_1(x)        
        x = self.conv_1_1(x)
        x = self.relu_1_2(x)
        x = self.norm_1_3(x)
        x = self.dequant_1(x)
        x = self.conv_1_4(x)
        x = self.relu_1_5(x)
        # continued...
</code></pre>
<p>I am attempting to quantize the model, then evaluate its performance on a dataset of interest:</p>
<pre><code class="lang-auto">    model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
    checkpoints_dir = '/spell/checkpoints'
    model.load_state_dict(
        torch.load(f"{checkpoints_dir}/model_50.pth", map_location=torch.device('cpu'))
    )
    model.eval()
    
    # NEW
    model = torch.quantization.prepare(model)
    print(f"Quantizing the model...")
    start_time = time.time()
    
    for i, (batch, segmap) in enumerate(dataloader):
        # batch = batch.cuda()
        # segmap = segmap.cuda()
        model(batch)

    model = torch.quantization.convert(model)
    print(f"Quantization done in {str(time.time() - start_time)} seconds.")

    print(f"Evaluating the model...")
    start_time = time.time()
    for i, (batch, segmap) in enumerate(dataloader):
        # batch = batch.cuda()
        # segmap = segmap.cuda()
        model(batch)
    
    print(f"Evaluation done in {str(time.time() - start_time)} seconds.")
</code></pre>
<p>The code, as written, fails with the following log output:</p>
<pre><code class="lang-auto">Loading the model...
Quantizing the model...
Quantization done in 84.06090354919434 seconds.
Evaluating the model...
Traceback (most recent call last):
  File "/spell/servers/eval_quantized.py", line 277, in &lt;module&gt;
    main()
  File "/spell/servers/eval_quantized.py", line 269, in main
    model(batch)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/spell/servers/eval_quantized.py", line 178, in forward
    x = self.conv_1_4(x)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/quantized/modules/conv.py", line 215, in forward
    self.dilation, self.groups, self.scale, self.zero_point)
RuntimeError: Could not run 'quantized::conv2d' with arguments from the 'CPUTensorId' backend. 'quantized::conv2d' is only available for these backends: [QuantizedCPUTensorId].
</code></pre>
<p>This implies that <code>x = self.conv_1_4(x)</code> is wrong, because <code>conv_1_4</code> is quantized. But I don’t understand <em>why</em> <code>conv_1_4</code> is quantized, though.</p>
<p>(apologies for posting an incomplete question–hit the wrong button)</p>

Vasiliy_Kuznetsov: <aside class="quote no-group" data-username="ResidentMario" data-post="1" data-topic="103311">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/residentmario/48/20188_2.png" class="avatar"> ResidentMario:</div>
<blockquote>
<pre><code class="lang-auto">        x = self.dequant_1(x)
        x = self.conv_1_4(x)
</code></pre>
</blockquote>
</aside>
<p>If the intent is to dequantize and want to do <code>conv_1_4</code> in fp32, then the problem is that by default, the quantization APIs quantize all convolutions in the model.  The workaround would be to disable quantization for <code>conv_1_4</code>, like this:</p>
<pre><code class="lang-auto">model.qconfig = ...
# disable quant for a specific layer
model.conv_1_4.qconfig = None
# continue with the quantization APIs, conv_1_4 will not be quantized
</code></pre>

ResidentMario: <blockquote>
<p>by default, the quantization APIs quantize all convolutions in the model</p>
</blockquote>
<p>Ah! This was not obvious from reading the documentation; the <a href="https://pytorch.org/docs/stable/quantization.html" rel="noopener nofollow ugc">Quantization quickstart</a> does not mention this behavior, nor, as best I can tell, does the <a href="https://pytorch.org/tutorials/advanced/static_quantization_tutorial.html" rel="noopener nofollow ugc">Static quantization tutorial</a>.</p>
<p>May I suggest making a docs task for this? Sorry, I know this is my second one in as many days. <img src="https://discuss.pytorch.org/images/emoji/apple/smile.png?v=12" title=":smile:" class="emoji" alt=":smile:" loading="lazy" width="20" height="20"></p>