Usage of tensor attributes in FX quantization

kenny1: <p>Assuming that a custom module is added to an original model a structure: `Linear → GELU → Linear’ using FX graph manipulation the GELU which doesn’t get quantized via FX mode is replaced by a custom module that works on int arithmetics to perform some approximations. I would like to know how to expose the quantized tensor as input to this custom module.</p>
<p>I tried setting the following approaches for the custom module</p>
<ul>
<li>Setting the backend_config for qint8 input and o/p</li>
<li>Setting the qconfig for qint8 for activation and weight</li>
<li>Manual Insertion of Stubs in the custom module</li>
</ul>
<p>However, the problem seems to be clear that these tensor attributes were indeed getting run on the dequantized tensor o/p of the previous linear layer as evident from the converted model.</p>
<p>GraphModule(<br>
(linear1): QuantizedLinear(in_features=256, out_features=1024, scale=1.0, zero_point=0, qscheme=torch.per_channel_affine)<br>
(gelu): Module()<br>
(linear2): QuantizedLinear(in_features=1024, out_features=256, scale=1.0, zero_point=0, qscheme=torch.per_channel_affine)<br>
)<br>
def forward(self, x):<br>
linear1_input_scale_0 = self.linear1_input_scale_0<br>
linear1_input_zero_point_0 = self.linear1_input_zero_point_0<br>
quantize_per_tensor = torch.quantize_per_tensor(x, linear1_input_scale_0, linear1_input_zero_point_0, torch.quint8);  x = linear1_input_scale_0 = linear1_input_zero_point_0 = None<br>
linear1 = self.linear1(quantize_per_tensor);  quantize_per_tensor = None<br>
dequantize_1 = linear1.dequantize();  linear1 = None<br>
gelu__c = self.gelu._c<br>
forward = gelu__c.forward(dequantize_1);  gelu__c = dequantize_1 = None<br>
gelu_scale_0 = self.gelu_scale_0<br>
gelu_zero_point_0 = self.gelu_zero_point_0<br>
quantize_per_tensor_2 = torch.quantize_per_tensor(forward, gelu_scale_0, gelu_zero_point_0, torch.quint8);  forward = gelu_scale_0 = gelu_zero_point_0 = None<br>
linear2 = self.linear2(quantize_per_tensor_2);  quantize_per_tensor_2 = None<br>
dequantize_3 = linear2.dequantize();  linear2 = None<br>
return dequantize_3</p>
<p>This indeed then results in the issue when performing inference of the converted model: “Could not run ‘aten::q_scale’'or aten::int_repr with arguments from the CPU backend”. Is there a way to perform some integer arithemetics on the quantized tensor such that FX quantization is compatible with this example custom module say:</p>
<p>class customGELU(nn.Module):<br>
def <strong>init</strong>(self):<br>
super(customGELU, self).<strong>init</strong>()<br>
def forward(self, x):<br>
q_val = x.int_repr()<br>
q_appr = q_val &gt;&gt; 2<br>
q_appr_float = (q_appr.float() - x.q_zero_point()) * x.q_scale()<br>
return torch.quantize_per_tensor(q_appr_float, x.q_scale(), x.q_zero_point(), x.dtype)</p>

jerryzh168: <p>fx also have a custom module quantization as well, maybe you can check out: <a href="https://github.com/pytorch/pytorch/blob/69e125a7e9b3ad50d993323f7a9dc2ee3edc6971/test/quantization/fx/test_quantize_fx.py#L3140" class="inline-onebox" rel="noopener nofollow ugc">pytorch/test/quantization/fx/test_quantize_fx.py at 69e125a7e9b3ad50d993323f7a9dc2ee3edc6971 · pytorch/pytorch · GitHub</a></p>