Quantization Bug in Concatenation of Tensor

Arkham_B: <p>Using AO Quantization package in pytorch I am doing post training quantization. I have a concatenation layer which concatenates 2 Quantized tensors with their own scale and zero_point. Ideally when 2 tensors are concatenated a new scaling and zero_point should be calculated implicitly due to change in overall min and max values. However in pytorch, scale and zero_point for the concatenated tensor are taken that of the first tensor in the concat list. No new value is calculated.</p>
<p>Please suggest solutions. I am sharing the code and the outputs for easier comprehension. As you can see tensor Y and tensor tmp as in code have same scale and zero_point.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F
class SimpleConvNet(nn.Module):
    def __init__(self):
        super(SimpleConvNet, self).__init__()
        self.quant = torch.quantization.QuantStub()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=3, stride=1, padding=1, bias=False)
        self.relu1 = nn.ReLU()
        self.conv2 = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=3, stride=1, padding=1, bias=False)
        self.relu2 = nn.ReLU()
        self.dequant = torch.quantization.DeQuantStub()
    
    def forward(self,x):
        x = self.quant(x)
        y = self.conv1(x)
        y = self.relu1(y)
        z = self.conv2(y)
        z = self.relu2(z)
        tmp = torch.cat([y,z],dim=1) #ISSUE HERE
        
        if y.dtype == torch.quint8:
            print('In Quantized Model')
            print('Tensor Y =&gt;',y,'\n\n')
            print('Tensor Z =&gt;',z,'\n\n')
            print('Concatenated Tensor =&gt;',tmp,'\n\n')
        tmp = self.dequant(tmp)
        return tmp

def main():
    device = 'cpu'#torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = SimpleConvNet().to(device)
    model.qconfig = torch.ao.quantization.default_qconfig#get_default_qconfig('x86')
    model = torch.ao.quantization.prepare(model)
    torch.manual_seed(0)
    inp = torch.rand((1,1,2,2))*20000-20
    inp = inp.to(device)
    output = model(inp)
    model = torch.ao.quantization.convert(model)
    model(inp)
    

if __name__ == "__main__":
    main()
</code></pre>
<p><strong>PRINT OUTPUT</strong><br>
In Quantized Model<br>
Tensor Y =&gt; tensor([[[[  85.7967, 1816.0299],<br>
[ 958.0630,  486.1812]]]], size=(1, 1, 2, 2), dtype=torch.quint8,<br>
quantization_scheme=torch.per_tensor_affine, scale=14.299448013305664,<br>
zero_point=0)</p>
<p>Tensor Z =&gt; tensor([[[[  0.0000, 293.9109],<br>
[ 92.2744,   0.0000]]]], size=(1, 1, 2, 2), dtype=torch.quint8,<br>
quantization_scheme=torch.per_tensor_affine, scale=3.417569160461426,<br>
zero_point=40)</p>
<p>Concatenated Tensor =&gt; tensor([[[[  85.7967, 1816.0299],<br>
[ 958.0630,  486.1812]],<br>
[[   0.0000,  300.2884],<br>
[  85.7967,    0.0000]]]], size=(1, 2, 2, 2), dtype=torch.quint8,<br>
quantization_scheme=torch.per_tensor_affine, scale=14.299448013305664,<br>
zero_point=0)</p>

HDCharles: <aside class="quote no-group" data-username="Arkham_B" data-post="1" data-topic="197505">
<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/arkham_b/48/67852_2.png" class="avatar"> Arkham_B:</div>
<blockquote>
<p>Please suggest solutions. I am sharing the code and the outputs for easier comprehension. As you can see tensor Y and tensor tmp as in code have same scale and zero_point.</p>
<pre><code class="lang-auto">
</code></pre>
</blockquote>
</aside>
<p>the intended solution is to use float_functional modules to do the cat operation which then should get quantized correctly. Incidentally this is also how you’d add/mul…etc.</p>
<p>in general eager mode quantization can’t handle functional operations, they have to be converted to modules so they can be interact with.</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/main/torch/ao/nn/quantized/modules/functional_modules.py#L9">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/main/torch/ao/nn/quantized/modules/functional_modules.py#L9" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/main/torch/ao/nn/quantized/modules/functional_modules.py#L9" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/main/torch/ao/nn/quantized/modules/functional_modules.py#L9</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="1" style="counter-reset: li-counter 0 ;">
          <li>from typing import List</li>
          <li></li>
          <li>import torch</li>
          <li>from torch import Tensor</li>
          <li>from torch._ops import ops</li>
          <li></li>
          <li>__all__ = ['FloatFunctional', 'FXFloatFunctional', 'QFunctional']</li>
          <li></li>
          <li class="selected">class FloatFunctional(torch.nn.Module):</li>
          <li>    r"""State collector class for float operations.</li>
          <li></li>
          <li>    The instance of this class can be used instead of the ``torch.`` prefix for</li>
          <li>    some operations. See example usage below.</li>
          <li></li>
          <li>    .. note::</li>
          <li></li>
          <li>        This class does not provide a ``forward`` hook. Instead, you must use</li>
          <li>        one of the underlying functions (e.g. ``add``).</li>
          <li></li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>in specific this method for cat:</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/main/torch/ao/nn/quantized/modules/functional_modules.py#L70">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/main/torch/ao/nn/quantized/modules/functional_modules.py#L70" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/main/torch/ao/nn/quantized/modules/functional_modules.py#L70" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/main/torch/ao/nn/quantized/modules/functional_modules.py#L70</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="60" style="counter-reset: li-counter 59 ;">
          <li>    return r</li>
          <li></li>
          <li>r"""Operation equivalent to ``torch.mul(Tensor, float)``"""</li>
          <li>def mul_scalar(self, x: Tensor, y: float) -&gt; Tensor:</li>
          <li>    r = torch.mul(x, y)</li>
          <li>    # Note: this operation is not observed because the observation is not</li>
          <li>    # needed for the quantized op.</li>
          <li>    return r</li>
          <li></li>
          <li>r"""Operation equivalent to ``torch.cat``"""</li>
          <li class="selected">def cat(self, x: List[Tensor], dim: int = 0) -&gt; Tensor:</li>
          <li>    r = torch.cat(x, dim=dim)</li>
          <li>    r = self.activation_post_process(r)</li>
          <li>    return r</li>
          <li></li>
          <li>r"""Operation equivalent to ``relu(torch.add(x,y))``"""</li>
          <li>def add_relu(self, x: Tensor, y: Tensor) -&gt; Tensor:</li>
          <li>    r = torch.add(x, y)</li>
          <li>    r = torch.nn.functional.relu(r)</li>
          <li>    r = self.activation_post_process(r)</li>
          <li>    return r</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>