Static Quantization of UNet

username: Surya_J
cooked: <p>Hi,</p>
<p>I am trying to quantize a UNet model using builtin static quantization functions.</p>
<ul>
<li>Pytorch CPU version 1.9.1</li>
<li>Ubuntu 20.04 LTS (conda env)</li>
</ul>
<p>The model itself is referenced from <a href="https://blog.jovian.ai/semantic-segmentation-in-self-driving-cars-3cb89aa08e9b" rel="noopener nofollow ugc">here</a>. I modified the model  as follows (showing the quantization parts alone) :</p>
<pre><code class="lang-auto">class UNet(nn.Module):
    def __init__(self, num_classes, quantize=False):        
        super(UNet, self).__init__()
        self.num_classes = num_classes
        
        """ QUANTIZED VERSION ADDITIONS """
        self.quantize = quantize
        self.quant = torch.quantization.QuantStub()
        self.dequant = torch.quantization.DeQuantStub()        

    def forward(self, X):
        # Outputs are dequantized
        if self.quantize == True:
            output_out = self.dequant(output_out)
        
        # pass through other layers
        
        # Outputs are dequantized
        if self.quantize == True:
            output_out = self.dequant(output_out)
        return output_out
</code></pre>
<p>The quantization function is as follows:</p>
<pre><code class="lang-auto">def quantizeUNet(model, device, dataLoader, use_fbgemm=False):
    model.to(device)    
    model.eval()    
    modules_to_fuse =  [['contracting_11.0', 'contracting_11.2'],
                        ['contracting_11.3', 'contracting_11.5'],
                        ['contracting_21.0', 'contracting_21.2'],
                        ['contracting_21.3', 'contracting_21.5'],
                        ['contracting_31.0', 'contracting_31.2'],
                        ['contracting_31.3', 'contracting_31.5'],
                        ['contracting_41.0', 'contracting_41.2'],
                        ['contracting_41.3', 'contracting_41.5'],
                        ['middle.0', 'middle.2'],
                        ['middle.3', 'middle.5'],
                        ['expansive_12.0', 'expansive_12.2'],
                        ['expansive_12.3', 'expansive_12.5'],
                        ['expansive_22.0', 'expansive_22.2'],
                        ['expansive_22.3', 'expansive_22.5'],
                        ['expansive_32.0', 'expansive_32.2'],
                        ['expansive_32.3', 'expansive_32.5'],
                        ['expansive_42.0', 'expansive_42.2'],
                        ['expansive_42.3', 'expansive_42.5']]
    #print(modules_to_fuse)
    model = torch.quantization.fuse_modules(model, modules_to_fuse)
    if use_fbgemm == True:
        model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
    else:
        model.qconfig = torch.quantization.default_qconfig
    torch.quantization.prepare(model, inplace=True)
    
    ## Calibrate Quantization parameters on input dataset
    print('Calibrating Quantization parameters on input dataset ...')
    model.eval()
    with torch.no_grad():
        for data, target in dataLoader:
            model(data)
    torch.quantization.convert(model, inplace=True)
    print('### Static Quantization complete ###')
    return model
</code></pre>
<p>During inference, the output tensor (shape [1, 10, 256, 256]) contains 0s only.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/0/10b14f1532465a3bc9eeb0b044265e746da513de.png" data-download-href="https://discuss.pytorch.org/uploads/default/10b14f1532465a3bc9eeb0b044265e746da513de" title="pytorch error"><img src="https://discuss.pytorch.org/uploads/default/original/3X/1/0/10b14f1532465a3bc9eeb0b044265e746da513de.png" alt="pytorch error" data-base62-sha1="2nFwn4fWJ4z4tJMyAPfAsyQQxky" width="584" height="499" data-dominant-color="F5F5F5"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">pytorch error</span><span class="informations">621×531 22.2 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I expected the output to have probabilities for each class (10 classes in total). But its essentially zero matrix. Is there something I’m missing? How to do static quantization of the model correctly?</p>

username: Vasiliy_Kuznetsov
cooked: <aside class="quote no-group" data-username="Surya_J" data-post="1" data-topic="134597">
<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/surya_j/48/43123_2.png" class="avatar"> Surya_J:</div>
<blockquote>
<pre><code class="lang-auto">    def forward(self, X):
        # Outputs are dequantized
        if self.quantize == True:
            output_out = self.dequant(output_out)
        
        # pass through other layers
        
        # Outputs are dequantized
        if self.quantize == True:
            output_out = self.dequant(output_out)
        return output_out
</code></pre>
</blockquote>
</aside>
<p>It might be a typo, but it should be something like</p>
<pre><code class="lang-auto"># quantize inputs (currently your code doesn't have this in what you pasted)
# run quantized model
# dequantize outputs
</code></pre>
<p>Other than that, your code looks right, as long as your calibration dataset is representative of your inference dataset.  Are you sure the expected output from your input is not close to a matrix of zeros?  Do you get the same thing for other outputs?</p>
<p>You could also try using <a href="https://pytorch.org/tutorials/prototype/numeric_suite_tutorial.html" class="inline-onebox" rel="noopener nofollow ugc">PyTorch Numeric Suite Tutorial — PyTorch Tutorials 2.1.1+cu121 documentation</a> to see if you can bisect the difference to a specific layer.</p>

username: Surya_J
cooked: <p>Hi <a class="mention" href="/u/vasiliy_kuznetsov">@Vasiliy_Kuznetsov</a>,</p>
<p>Thanks for the reply. Yes, I pasted the wrong code here. This is my actual code is :</p>
<pre><code class="lang-auto">def forward(self, X):
    # Input are quantized
    if self.quantize == True:
        X = self.quant(X)
</code></pre>
<p>The output is zero for the entire test set (I’m using a subset of the CityScapes dataset). The un-quantized model gives floating point output and the predictions are good. So, I assume there’s something missing when I quantize the model. I’ll try debugging using the link you shared. Thanks for the help.</p>

username: Surya_J
cooked: <p><strong>Update</strong></p>
<ul>
<li>
<p>It seems that ConvTranspose2d is not yet supported for quantization. Hence, you have to dequantize the output before passing through each of the unsupported layers, which is slower than the original float model in my case.  <a href="https://discuss.pytorch.org/t/is-there-a-way-to-quantize-conv-transpose2d-layer/96590/31">Related Forum post</a></p>
</li>
<li>
<p>I guess it;s better to look for models which contain only supported layers in case of static quantization.</p>
</li>
</ul>

username: jerryzh168
cooked: <p>I think convtranspose is supported: <a href="https://github.com/pytorch/pytorch/blob/master/torch/ao/quantization/quantization_mappings.py#L43" class="inline-onebox" rel="noopener nofollow ugc">pytorch/quantization_mappings.py at master · pytorch/pytorch · GitHub</a></p>

username: Surya_J
cooked: <p>Hi <a class="mention" href="/u/jerryzh168">@jerryzh168</a>,</p>
<p>Thanks for reaching out.</p>
<p>Based on the forum post, I jumped to that conclusion that <code>ConvTranspose2d</code> is not supported yet. Will look into the model and check for the problem.</p>