Updated math model to target int8 x int8 kernels.

math_model.py

@@ -47,11 +47,12 @@ class QuantLinear(nn.Module):
         #  - multiply this sum with every weight zero-point (e.g., `torch.sum(quant_input, dim=-1) * self.weight_zp`
         #  - Subtract from previous output (e.g., `quant_output -= torch.sum(quant_input, dim=-1) * self.weight_zp`)
         #  - All other code is just to make sure the broadcasting semantics work correctly
-        quant_weight = quantize(self.linear.weight, self.weight_scale, self.weight_zp, is_asym=True).to(torch.uint8)
+        weight_zp_int8 = (self.weight_zp - 128).to(torch.int8).to(torch.float32)
+        quant_weight = quantize(self.linear.weight, self.weight_scale, weight_zp_int8, is_asym=False).to(torch.int8)
         fused_input_scale = self.input_scale / self.mul_factor # Fuse SmoothQuant and input scales, can be computed offline
         quant_input = quantize(x, fused_input_scale, self.input_zp, is_asym=False).to(torch.int8)
         quant_output = torch.nn.functional.linear(quant_input.to(torch.float32), quant_weight.to(torch.float32), None).to(torch.int32) # Convert inputs to FP32 to avoid F.linear quantizing the output to int8
-        correction = torch.sum(quant_input, dim=-1, keepdim=True).to(torch.int32) * (-self.weight_zp).to(torch.uint8).view([1]*(quant_input.ndim-1) + [self.weight_zp.nelement()]) # Correct for weight zero-point
+        correction = torch.sum(quant_input, dim=-1, keepdim=True).to(torch.int32) * (-weight_zp_int8).to(torch.int8).view([1]*(quant_input.ndim-1) + [self.weight_zp.nelement()]) # Correct for weight zero-point
         quant_output = quant_output + correction
         output = dequantize(quant_output, (self.weight_scale * self.input_scale).view([1]*(quant_output.ndim-1) + [(self.weight_scale * self.input_scale).nelement()]), 0.0)
         output += self.linear.bias
@@ -103,15 +104,16 @@ class QuantConv2d(nn.Module):
         #  - multiply this sum with every weight zero-point (e.g., `sum * self.weight_zp`
         #  - Subtract from previous output (e.g., `quant_output -= sum * self.weight_zp`)
         #  - All other code is just to make sure the broadcasting semantics work correctly
-        quant_weight = quantize(self.conv2d.weight, self.weight_scale, self.weight_zp, is_asym=True).to(torch.uint8)
+        weight_zp_int8 = (self.weight_zp - 128).to(torch.int8).to(torch.float32)
+        quant_weight = quantize(self.conv2d.weight, self.weight_scale, weight_zp_int8, is_asym=False).to(torch.int8)
         b_shape = list(quant_weight.shape) # Used for weight zero-point correction
         b_shape[0] = 1 # Used for weight zero-point correction
-        weight_cat = torch.ones((1,1,1,1)).broadcast_to(b_shape).to(torch.uint8) # Used for weight zero-point correction
-        quant_weight = torch.cat((quant_weight,weight_cat),dim=0).to(torch.uint8) # Create extra output channel, used for weight zero-point correction
+        weight_cat = torch.ones((1,1,1,1)).broadcast_to(b_shape).to(torch.int8) # Used for weight zero-point correction
+        quant_weight = torch.cat((quant_weight,weight_cat),dim=0).to(torch.int8) # Create extra output channel, used for weight zero-point correction
         fused_input_scale = self.input_scale / self.mul_factor # Fuse SmoothQuant and input scales, can be computed offline
         quant_input = quantize(x, fused_input_scale, self.input_zp, is_asym=False).to(torch.int8)
         quant_output = torch.nn.functional.conv2d(quant_input.to(torch.float32), quant_weight.to(torch.float32), None).to(torch.int32) # Convert inputs to FP32 to avoid F.conv2d quantizing the output to int8
-        correction = quant_output[:,-1,:,:] * (-self.weight_zp).to(torch.uint8).view([1, self.weight_zp.nelement()] +  [1]*(quant_output.ndim-2)) # Correct zero-point for weight
+        correction = quant_output[:,-1,:,:] * (-weight_zp_int8).to(torch.int8).view([1, self.weight_zp.nelement()] +  [1]*(quant_output.ndim-2)) # Correct zero-point for weight
         quant_output = quant_output[:,:-1,:,:] + correction
         output = dequantize(quant_output, (self.weight_scale * self.input_scale).view([1, (self.weight_scale * self.input_scale).nelement()] + [1]*(quant_output.ndim-2)), 0.0)
         output += self.conv2d.bias.view([1, self.conv2d.bias.nelement()] + [1]*(quant_output.ndim-2))

test_quant_conv2d.py

@@ -1,23 +1,28 @@
 import torch
+import torch.nn as nn
 from math_model import QuantConv2d
 
 torch.manual_seed(0)
 
 batch_size = 1
-out_ch = 8
-in_ch = 4
+out_ch = 128
+in_ch = 64
 k = 3
 h = 5
 w = 5
 
+i = 2*torch.rand((batch_size,in_ch,h,w)) - 1.
+l = nn.Conv2d(in_ch, out_ch, k, bias=True)
+
 quant_params = {
     'smoothquant_mul': torch.rand((in_ch,)),
     'smoothquant_mul_shape': (1,in_ch,1,1),
     'weight_scale': torch.rand((out_ch,)),
+    'weight_scale': torch.max(torch.abs(torch.flatten(l.weight, start_dim=1)), dim=1).values / 128.,
     'weight_scale_shape': (out_ch,1,1,1),
-    'weight_zp': torch.randint(-255, 0, (out_ch,)),
+    'weight_zp': torch.clamp(torch.round((torch.mean((l.weight), dim=(1,2,3))) * (128 / torch.max(torch.abs(torch.flatten(l.weight, start_dim=1)), dim=1).values)) + 128, 0, 255),
     'weight_zp_shape': (out_ch,1,1,1),
-    'input_scale': torch.rand((1,)),
+    'input_scale': torch.max(torch.abs(i)) / 128.,
     'input_scale_shape': tuple(),
     'input_zp': torch.zeros((1,)),
     'input_zp_shape': tuple(),
@@ -25,10 +30,10 @@ quant_params = {
 
 print(quant_params)
 
-l = QuantConv2d(in_ch, out_ch, k, quant_params)
-i = torch.rand((batch_size,in_ch,h,w))
-o_qdq = l(i)
-o_qop = l(i, qop=True)
+ql = QuantConv2d(in_ch, out_ch, k, quant_params)
+ql.conv2d.load_state_dict(l.state_dict())
+o_qdq = ql(i)
+o_qop = ql(i, qop=True)
 print(o_qdq.shape)
 print(o_qop.shape)
 print(o_qdq - o_qop)

test_quant_linear.py

@@ -1,4 +1,5 @@
 import torch
+import torch.nn as nn
 from math_model import QuantLinear
 
 torch.manual_seed(0)
@@ -7,14 +8,17 @@ batch_size = 1
 out_ch = 128
 in_ch = 64
 
+i = 2*torch.rand((batch_size,in_ch)) - 1.
+l = nn.Linear(in_ch, out_ch, bias=True)
+
 quant_params = {
     'smoothquant_mul': torch.rand((in_ch,)),
     'smoothquant_mul_shape': (1,in_ch),
-    'weight_scale': torch.rand((out_ch,)),
+    'weight_scale': torch.max(torch.abs(l.weight), dim=1).values / 128.,
     'weight_scale_shape': (out_ch,1),
-    'weight_zp': torch.randint(-255, 0, (out_ch,)),
+    'weight_zp': torch.clamp(torch.round((torch.mean((l.weight), dim=1)) * (128 / torch.max(torch.abs(l.weight), dim=1).values)) + 128, 0, 255),
     'weight_zp_shape': (out_ch,1),
-    'input_scale': torch.rand((1,)),
+    'input_scale': torch.max(torch.abs(i)) / 128.,
     'input_scale_shape': tuple(),
     'input_zp': torch.zeros((1,)),
     'input_zp_shape': tuple(),
@@ -22,10 +26,10 @@ quant_params = {
 
 print(quant_params)
 
-l = QuantLinear(in_ch, out_ch, quant_params)
-i = torch.rand((batch_size,in_ch))
-o_qdq = l(i)
-o_qop = l(i, qop=True)
+ql = QuantLinear(in_ch, out_ch, quant_params)
+ql.linear.load_state_dict(l.state_dict())
+o_qdq = ql(i)
+o_qop = ql(i, qop=True)
 print(o_qdq.shape)
 print(o_qop.shape)
 print(o_qdq - o_qop)