Is there an alternative to do batched matrix multiplication on Quantized Tensors?

harshsaini: <p>Hi,</p>
<p>I am trying to do post training static quantization, however, I am running into issues where certain operations are not defined for <code>QuantizedCPUTensorId</code>.</p>
<p>Minimal reproducible example:</p>
<pre><code class="lang-auto">&gt;&gt;&gt; import torch
&gt;&gt;&gt; 
&gt;&gt;&gt; A = torch.Tensor([[2,2], [3,3]]).unsqueeze(0)
&gt;&gt;&gt; B = torch.Tensor([[2,3], [2,3]]).unsqueeze(0)
&gt;&gt;&gt; scale, zero_point, dtype = 1.0, 2, torch.qint8
&gt;&gt;&gt; qA = torch.quantize_per_tensor(A, scale, zero_point, dtype)
&gt;&gt;&gt; qB = torch.quantize_per_tensor(B, scale, zero_point, dtype)
&gt;&gt;&gt; torch.matmul(A,B)
tensor([[[ 8., 12.],
         [12., 18.]]])
&gt;&gt;&gt; torch.matmul(qA,qB)
Traceback (most recent call last):
  File "&lt;stdin&gt;", line 1, in &lt;module&gt;
RuntimeError: Could not run 'aten::bmm' with arguments from the 'QuantizedCPUTensorId' backend. 'aten::bmm' is only available for these backends: [CPUTensorId, VariableTensorId].
</code></pre>
<p>Are there alternatives to accomplishing the same?<br>
I know there are certain operations that are defined here: <a href="https://pytorch.org/docs/stable/quantization.html#floatfunctional" rel="nofollow noopener">https://pytorch.org/docs/stable/quantization.html#floatfunctional</a> but what would be the optimal way?</p>

supriyar: <p>If possible try using <code>nn.Linear</code> instead of <code>aten::bmm</code>.</p>
<p>Currently the only way is to implement the quantized operator for <code>aten::bmm</code>.<br>
One easy way could be by implementing the <code>quantized::linear</code> operator by looping over the batch dimension. We will be looking into implementing this operator in the future.</p>

harshsaini: <p>Hi <a class="mention" href="/u/supriyar">@supriyar</a>, thanks for the response.</p>
<p>Yes, I had thought about that but wouldn’t that operation be suboptimal? However, if there is no alternative, I guess it would have to be so for now.</p>

harshsaini: <p>Seems like <a href="https://pytorch.org/docs/stable/quantization.html#torch.nn.quantized.functional.linear" rel="nofollow noopener">https://pytorch.org/docs/stable/quantization.html#torch.nn.quantized.functional.linear</a> is not a viable option. It requires the input tensor to be unsigned, however, the operation explicitly is between two tensors that are <code>qint8</code>.</p>
<pre><code class="lang-auto">&gt;&gt;&gt; torch.nn.quantized.functional.linear(qA[0,], qB[0,])
RuntimeError: expected scalar type QUInt8 but found QInt8
</code></pre>

jerryzh168: <p>do you need both of the inputs to be <code>qint8</code>? If you change <code>qA</code> to be <code>quint8</code> it would work</p>

itayhubara: <p>Any update about this one? Are you going to support it in the near future?</p>

ignatius: <p>I’ve came across with the following code. As pointed out I perform the matrix multiplication with nn.Linear.</p>
<p>What do you think?</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

class BatchedMatMul(nn.Module):
    def __init__(self):
        super().__init__()
        self.quant = torch.quantization.QuantStub()
        self.linear = nn.Linear(3,3, bias=False)
        self.dequant = torch.quantization.DeQuantStub()

    def forward(self, input1, input2):
        y = []
        for b in range(input1.shape[0]):
          print(f"Linear's type: {type(self.linear)}")
          print(f"Linear's weigth type: {type(self.linear.weight)}")
          if isinstance(self.linear.weight, nn.Parameter):
            self.linear.weight.requires_grad = False
            self.linear.weight.copy_ (self.quant(input1[b]))
            y.append(self.linear(self.quant(input2[b])))
          else:
            scale = self.linear.weight().q_per_channel_scales()
            zero_point = self.linear.weight().q_per_channel_zero_points()
            w = torch.quantize_per_channel(input1[b], scale, zero_point, 1, torch.qint8)
            self.linear.set_weight_bias(w, b=None)
            y.append(self.linear(self.quant(input2[b])))
          
        return self.dequant(torch.stack(y))

print("Cronstruct model...")
matmul = BatchedMatMul()
print("Cronstruct model... [OK]")

matmul.eval()
print("Running FP32 inference...")
inp = torch.ones(3, 3).repeat(2,1,1)
y = matmul(2*inp, inp)
print("FP32 output...")
print(y)
print("Running FP32 inference... [OK]")

print("Quantizing...")
matmul.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
matmul_prepared = torch.quantization.prepare(matmul)
matmul_prepared(2*inp, inp)
model_int8 = torch.quantization.convert(matmul_prepared)
print("Quantizing... [OK]")
print("Running INT8 inference...")
y = model_int8.forward(2*inp, inp)
print("Int8 Output")
print(y)
print("Running INT8 inference..[OK]")
</code></pre>
<p>Output:</p>
<pre><code class="lang-auto">Cronstruct model...
Cronstruct model... [OK]
Running FP32 inference...
Linear's type: &lt;class 'torch.nn.modules.linear.Linear'&gt;
Linear's weigth type: &lt;class 'torch.nn.parameter.Parameter'&gt;
Linear's type: &lt;class 'torch.nn.modules.linear.Linear'&gt;
Linear's weigth type: &lt;class 'torch.nn.parameter.Parameter'&gt;
FP32 output...
tensor([[[6., 6., 6.],
         [6., 6., 6.],
         [6., 6., 6.]],

        [[6., 6., 6.],
         [6., 6., 6.],
         [6., 6., 6.]]])
Running FP32 inference... [OK]
Quantizing...
Linear's type: &lt;class 'torch.nn.modules.linear.Linear'&gt;
Linear's weigth type: &lt;class 'torch.nn.parameter.Parameter'&gt;
Linear's type: &lt;class 'torch.nn.modules.linear.Linear'&gt;
Linear's weigth type: &lt;class 'torch.nn.parameter.Parameter'&gt;
Quantizing... [OK]
Running INT8 inference...
Linear's type: &lt;class 'torch.nn.quantized.modules.linear.Linear'&gt;
Linear's weigth type: &lt;class 'method'&gt;
Linear's type: &lt;class 'torch.nn.quantized.modules.linear.Linear'&gt;
Linear's weigth type: &lt;class 'method'&gt;
Int8 Output
tensor([[[5.9695, 5.9695, 5.9695],
         [5.9695, 5.9695, 5.9695],
         [5.9695, 5.9695, 5.9695]],

        [[5.9695, 5.9695, 5.9695],
         [5.9695, 5.9695, 5.9695],
         [5.9695, 5.9695, 5.9695]]])
Running INT8 inference..[OK]
/usr/local/lib/python3.6/dist-packages/torch/quantization/observer.py:121: UserWarning: Please use quant_min and quant_max to specify the range for observers.                     reduce_range will be deprecated in a future release of PyTorch.
  reduce_range will be deprecated in a future release of PyTorch."
</code></pre>

jerryzh168: <p>currently we only support quint8 for activations and qint8 for weight I think.</p>
<p>Currently we do not have plans for supporting bmm, one workaround is to put DeQuantStub and QuantStub around bmm op to skip quantizing it.</p>