Change Linear layer weights in a quantized model

username: ignatius
cooked: <p>Since matrix multiplication is not supported for model quantization I’m performing it with a nn.Linear layer which I change its weigths in every forward pass.</p>
<p>This approach works well for the FP32 model but it crashes when the model is quantized. The issue is that, when the model is converted to int8, the following lines of code are not valid</p>
<pre><code class="lang-auto">self.linear.weight.requires_grad = False
self.linear.weight.copy_ (input1[b])
</code></pre>
<p>because in the converted model <code>self.linear.weight</code> is not a <code>torch.nn.Paramater</code> but a method which returns a Tensor</p>
<p>Any workarround on this?</p>
<p><strong>FULL CODE</strong></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)}")
          self.linear.weight.requires_grad = False
          self.linear.weight.copy_ (self.quant(input1[b]))
          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(inp, inp)

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(inp, inp)

model_int8 = torch.quantization.convert(matmul_prepared)

print("Quantizing... [OK]")

print("Running INT8 inference...")

y = model_int8(inp, inp)

print(y)

print("Running INT8 inference..[OK]")
</code></pre>
<p><strong>OUTPUT</strong></p>
<pre><code class="lang-auto">Cronstruct model...
Cronstruct model... [OK]
Running FP32 inference...
Linear's weigth type: &lt;class 'torch.nn.parameter.Parameter'&gt;
Linear's weigth type: &lt;class 'torch.nn.parameter.Parameter'&gt;
tensor([[[3., 3., 3.],
         [3., 3., 3.],
         [3., 3., 3.]],

        [[3., 3., 3.],
         [3., 3., 3.],
         [3., 3., 3.]]])
Running FP32 inference... [OK]
Quantizing...
Linear's weigth type: &lt;class 'torch.nn.parameter.Parameter'&gt;
Linear's weigth type: &lt;class 'torch.nn.parameter.Parameter'&gt;
Quantizing... [OK]
Running INT8 inference...
Linear's weigth type: &lt;class 'method'&gt;
/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."
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
&lt;ipython-input-81-024fd82f94de&gt; in &lt;module&gt;()
     34 print("Quantizing... [OK]")
     35 print("Running INT8 inference...")
---&gt; 36 y = model_int8(inp, inp)
     37 print(y)
     38 print("Running INT8 inference..[OK]")

1 frames
&lt;ipython-input-81-024fd82f94de&gt; in forward(self, input1, input2)
     10         for b in range(input1.shape[0]):
     11           print(f"Linear's weigth type: {type(self.linear.weight)}")
---&gt; 12           self.linear.weight.requires_grad = False
     13           self.linear.weight.copy_ (input1[b])
     14           y.append(self.linear(input2[b]))

AttributeError: 'method' object has no attribute 'requires_grad'
</code></pre>

username: ignatius
cooked: <p>I come acrross with the following:</p>
<pre><code class="lang-auto">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:
            self.linear.set_weight_bias(self.quant(input1[b]), b=None)
            y.append(self.linear(self.quant(input2[b])))
          
        return self.dequant(torch.stack(y))
</code></pre>
<p><code>self.linear</code> has changeg from  <em>torch.nn.modules.linear.Linear</em> to <em>torch.nn.quantized.modules.linear.Linear</em> so their methods and attributes are different.</p>
<p>Nevertheless, this approach is still throwing an error because the quantized linear layer expects an signed integer as parameter but an unsigned is being given…</p>
<pre><code class="lang-auto">---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-112-629b4f9ebdc0&gt; in &lt;module&gt;()
     42 print("Quantizing... [OK]")
     43 print("Running INT8 inference...")
---&gt; 44 y = model_int8.forward(inp, inp)
     45 print(y)
     46 print("Running INT8 inference..[OK]")

2 frames
/usr/local/lib/python3.6/dist-packages/torch/nn/quantized/modules/linear.py in set_weight_bias(self, weight, bias)
     21     def set_weight_bias(self, weight: torch.Tensor, bias: Optional[torch.Tensor]) -&gt; None:
     22         if self.dtype == torch.qint8:
---&gt; 23             self._packed_params = torch.ops.quantized.linear_prepack(weight, bias)
     24         elif self.dtype == torch.float16:
     25             self._packed_params = torch.ops.quantized.linear_prepack_fp16(weight, bias)

RuntimeError: expected scalar type QInt8 but found QUInt8
</code></pre>

username: ignatius
cooked: <p>I think I’m close… 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><strong>OUT</strong></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>

username: jerryzh168
cooked: <p>this feels a bit hacky, I’m not sure if it would work or not, bmm is not supported right now, why not put dequantstub and quantstub around bmm op to avoid quantizing it?</p>

username: ignatius
cooked: <p>Yes, it’s quite hacky… indeed, it only works for dynamic quantization (quantizing only the weights)</p>
<p>I’m very interested in placing the <strong>quantsub</strong> and <strong>dequantsub</strong> to avoid quantization, could you please provide a pice of code?</p>
<p>Do you mean something like this:</p>
<pre><code class="lang-auto"># y and x Tensors previouslly created...
x = torch.quantization.DeQuantStub(x)
y = torch.quantization.DeQuantStub(y)
Y = torch.torch.bmm(y, x)
y = torch.quantization.QuantStub(y)
</code></pre>
<p>Thx you very much!</p>

username: jerryzh168
cooked: <p>yes that’s exactly what I meant, did you try this?</p>

username: ignatius
cooked: <p>well. since quantization is not yet available for GPU inference… it is not worth for me to try this out…</p>