calibration support for quantized flipout layers

Author: junliang-lin

bayesian_torch/layers/flipout_layers/conv_flipout.py

@@ -37,6 +37,8 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from ..base_variational_layer import BaseVariationalLayer_, get_kernel_size
+from torch.quantization.observer import HistogramObserver, PerChannelMinMaxObserver, MinMaxObserver
+from torch.quantization.qconfig import QConfig
 
 from torch.distributions.normal import Normal
 from torch.distributions.uniform import Uniform
@@ -136,6 +138,15 @@ def __init__(self,
             self.register_buffer('prior_bias_sigma', None, persistent=False)
 
         self.init_parameters()
+        self.quant_prepare=False
+    
+    def prepare(self):
+        self.qint_quant = nn.ModuleList([torch.quantization.QuantStub(
+                                         QConfig(weight=MinMaxObserver.with_args(dtype=torch.qint8, qscheme=torch.per_tensor_symmetric), activation=MinMaxObserver.with_args(dtype=torch.qint8,qscheme=torch.per_tensor_symmetric))) for _ in range(4)])
+        self.quint_quant = nn.ModuleList([torch.quantization.QuantStub(
+                                         QConfig(weight=MinMaxObserver.with_args(dtype=torch.quint8), activation=MinMaxObserver.with_args(dtype=torch.quint8))) for _ in range(8)])
+        self.dequant = torch.quantization.DeQuantStub()
+        self.quant_prepare=True
 
     def init_parameters(self):
         # prior values
@@ -303,6 +314,15 @@ def __init__(self,
             self.register_buffer('prior_bias_sigma', None, persistent=False)
 
         self.init_parameters()
+        self.quant_prepare=False
+    
+    def prepare(self):
+        self.qint_quant = nn.ModuleList([torch.quantization.QuantStub(
+                                         QConfig(weight=MinMaxObserver.with_args(dtype=torch.qint8, qscheme=torch.per_tensor_symmetric), activation=MinMaxObserver.with_args(dtype=torch.qint8,qscheme=torch.per_tensor_symmetric))) for _ in range(4)])
+        self.quint_quant = nn.ModuleList([torch.quantization.QuantStub(
+                                         QConfig(weight=MinMaxObserver.with_args(dtype=torch.quint8), activation=MinMaxObserver.with_args(dtype=torch.quint8))) for _ in range(8)])
+        self.dequant = torch.quantization.DeQuantStub()
+        self.quant_prepare=True
 
     def init_parameters(self):
         # prior values
@@ -365,18 +385,38 @@ def forward(self, x, return_kl=True):
                                       self.prior_bias_sigma)
 
         # perturbed feedforward
-        perturbed_outputs = F.conv2d(x * sign_input,
+        x_tmp = x * sign_input
+        perturbed_outputs_tmp = F.conv2d(x * sign_input,
                                      weight=delta_kernel,
                                      bias=bias,
                                      stride=self.stride,
                                      padding=self.padding,
                                      dilation=self.dilation,
-                                     groups=self.groups) * sign_output
+                                     groups=self.groups)
+        perturbed_outputs = perturbed_outputs_tmp * sign_output
+        out = outputs + perturbed_outputs
+
+        if self.quant_prepare:
+            # quint8 quantstub
+            input = self.quint_quant[0](input) # input
+            outputs = self.quint_quant[1](outputs) # output
+            sign_input = self.quint_quant[2](sign_input)
+            sign_output = self.quint_quant[3](sign_output)
+            x_tmp = self.quint_quant[4](x_tmp)
+            perturbed_outputs_tmp = self.quint_quant[5](perturbed_outputs_tmp) # output
+            perturbed_outputs = self.quint_quant[6](perturbed_outputs) # output
+            out = self.quint_quant[7](out) # output
+
+            # qint8 quantstub
+            sigma_weight = self.qint_quant[0](sigma_weight) # weight
+            mu_kernel = self.qint_quant[1](self.mu_kernel) # weight
+            eps_kernel = self.qint_quant[2](eps_kernel) # random variable
+            delta_kernel =self.qint_quant[3](delta_kernel) # multiply activation
 
         # returning outputs + perturbations
         if return_kl:
-            return outputs + perturbed_outputs, kl
-        return outputs + perturbed_outputs
+            return out, kl
+        return out
 
 
 class Conv3dFlipout(BaseVariationalLayer_):

bayesian_torch/layers/flipout_layers/linear_flipout.py

@@ -40,6 +40,8 @@
 from torch.distributions.normal import Normal
 from torch.distributions.uniform import Uniform
 from ..base_variational_layer import BaseVariationalLayer_
+from torch.quantization.observer import HistogramObserver, PerChannelMinMaxObserver, MinMaxObserver
+from torch.quantization.qconfig import QConfig
 
 __all__ = ["LinearFlipout"]
 
@@ -107,6 +109,15 @@ def __init__(self,
             self.register_buffer('eps_bias', None, persistent=False)
 
         self.init_parameters()
+        self.quant_prepare=False
+    
+    def prepare(self):
+        self.qint_quant = nn.ModuleList([torch.quantization.QuantStub(
+                                         QConfig(weight=MinMaxObserver.with_args(dtype=torch.qint8, qscheme=torch.per_tensor_symmetric), activation=MinMaxObserver.with_args(dtype=torch.qint8,qscheme=torch.per_tensor_symmetric))) for _ in range(4)])
+        self.quint_quant = nn.ModuleList([torch.quantization.QuantStub(
+                                         QConfig(weight=MinMaxObserver.with_args(dtype=torch.quint8), activation=MinMaxObserver.with_args(dtype=torch.quint8))) for _ in range(8)])
+        self.dequant = torch.quantization.DeQuantStub()
+        self.quant_prepare=True
 
     def init_parameters(self):
         # init prior mu
@@ -136,7 +147,9 @@ def forward(self, x, return_kl=True):
             return_kl = False
         # sampling delta_W
         sigma_weight = torch.log1p(torch.exp(self.rho_weight))
-        delta_weight = (sigma_weight * self.eps_weight.data.normal_())
+        eps_weight = self.eps_weight.data.normal_()
+        delta_weight = sigma_weight * eps_weight
+        # delta_weight = (sigma_weight * self.eps_weight.data.normal_())
 
         # get kl divergence
         if return_kl:
@@ -153,14 +166,32 @@ def forward(self, x, return_kl=True):
 
         # linear outputs
         outputs = F.linear(x, self.mu_weight, self.mu_bias)
-
         sign_input = x.clone().uniform_(-1, 1).sign()
         sign_output = outputs.clone().uniform_(-1, 1).sign()
-
-        perturbed_outputs = F.linear(x * sign_input, delta_weight,
-                                     bias) * sign_output
+        x_tmp = x * sign_input
+        perturbed_outputs_tmp = F.linear(x_tmp, delta_weight, bias)
+        perturbed_outputs = perturbed_outputs_tmp * sign_output
+        out = outputs + perturbed_outputs
+
+        if self.quant_prepare:
+            # quint8 quantstub
+            input = self.quint_quant[0](input) # input
+            outputs = self.quint_quant[1](outputs) # output
+            sign_input = self.quint_quant[2](sign_input)
+            sign_output = self.quint_quant[3](sign_output)
+            x_tmp = self.quint_quant[4](x_tmp)
+            perturbed_outputs_tmp = self.quint_quant[5](perturbed_outputs_tmp) # output
+            perturbed_outputs = self.quint_quant[6](perturbed_outputs) # output
+            out = self.quint_quant[7](out) # output
+
+            # qint8 quantstub
+            sigma_weight = self.qint_quant[0](sigma_weight) # weight
+            mu_weight = self.qint_quant[1](self.mu_weight) # weight
+            eps_weight = self.qint_quant[2](eps_weight) # random variable
+            delta_weight =self.qint_quant[3](delta_weight) # multiply activation
+            
 
         # returning outputs + perturbations
         if return_kl:
-            return outputs + perturbed_outputs, kl
-        return outputs + perturbed_outputs
+            return out, kl
+        return out

bayesian_torch/layers/flipout_layers/quantized_conv_flipout.py

@@ -284,6 +284,7 @@ def __init__(self,
         self.bn_eps = None
 
         self.is_dequant = False
+        self.quant_dict = None
 
     def get_scale_and_zero_point(self, x, upper_bound=100, target_range=255):
         """ An implementation for symmetric quantization
@@ -425,40 +426,67 @@ def forward(self, x, normal_scale=6/255, default_scale=0.1, default_zero_point=1
         if self.dnn_to_bnn_flag:
             return_kl = False
 
-        if x.dtype!=torch.quint8:
-            x = torch.quantize_per_tensor(x, default_scale, default_zero_point, torch.quint8)
-
-        bias = None
-        if self.bias:
-            bias = self.quantized_mu_bias
-
-        outputs = torch.nn.quantized.functional.conv2d(x, self.quantized_mu_weight, bias, self.stride, self.padding,
-                        self.dilation, self.groups, scale=default_scale, zero_point=default_zero_point) # input: quint8, weight: qint8, bias: fp32
-
-        # sampling perturbation signs
-        sign_input = torch.zeros(x.shape).uniform_(-1, 1).sign()
-        sign_output = torch.zeros(outputs.shape).uniform_(-1, 1).sign()
-        sign_input = torch.quantize_per_tensor(sign_input, default_scale, default_zero_point, torch.quint8)
-        sign_output = torch.quantize_per_tensor(sign_output, default_scale, default_zero_point, torch.quint8)
-
-        # getting perturbation weights
-        eps_kernel = torch.quantize_per_tensor(self.eps_kernel.data.normal_(), normal_scale, 0, torch.qint8)
-        new_scale = (self.quantized_sigma_weight.q_scale())*(eps_kernel.q_scale())
-        delta_kernel = torch.ops.quantized.mul(self.quantized_sigma_weight, eps_kernel, new_scale, 0)
-
         bias = None
         if self.bias:
-            eps_bias = self.eps_bias.data.normal_()
-            bias = (self.quantized_sigma_bias * eps_bias)
+            bias = self.quantized_mu_bias # TODO: check correctness
+
+        if self.quant_dict is not None:
+            # getting perturbation weights
+            eps_kernel = torch.quantize_per_tensor(self.eps_kernel.data.normal_(), self.quant_dict[0]['scale'], self.quant_dict[0]['zero_point'], torch.qint8)
+            delta_kernel = torch.ops.quantized.mul(self.quantized_sigma_weight, eps_kernel, self.quant_dict[1]['scale'], self.quant_dict[1]['zero_point'])
+
+            if x.dtype!=torch.quint8: # check if input has been quantized
+                x = torch.quantize_per_tensor(x, self.quant_dict[2]['scale'], self.quant_dict[2]['zero_point'], torch.quint8) # scale=0.1 by grid search; zero_point=128 for uint8 format
+
+            outputs = torch.nn.quantized.functional.conv2d(x, self.quantized_mu_weight, bias, self.stride, self.padding,
+                            self.dilation, self.groups, scale=self.quant_dict[3]['scale'], zero_point=self.quant_dict[3]['zero_point']) # input: quint8, weight: qint8, bias: fp32
+
+            # sampling perturbation signs
+            sign_input = torch.zeros(x.shape).uniform_(-1, 1).sign()
+            sign_output = torch.zeros(outputs.shape).uniform_(-1, 1).sign()
+            sign_input = torch.quantize_per_tensor(sign_input, self.quant_dict[4]['scale'], self.quant_dict[4]['zero_point'], torch.quint8)
+            sign_output = torch.quantize_per_tensor(sign_output, self.quant_dict[5]['scale'], self.quant_dict[5]['zero_point'], torch.quint8)
+            
+            # perturbed feedforward
+            x = torch.ops.quantized.mul(x, sign_input, self.quant_dict[6]['scale'], self.quant_dict[6]['zero_point'])
+            perturbed_outputs = torch.nn.quantized.functional.conv2d(x,
+                                weight=delta_kernel, bias=bias, stride=self.stride, padding=self.padding,
+                                dilation=self.dilation, groups=self.groups, scale=self.quant_dict[7]['scale'], zero_point=self.quant_dict[7]['zero_point'])
+            perturbed_outputs = torch.ops.quantized.mul(perturbed_outputs, sign_output, self.quant_dict[8]['scale'], self.quant_dict[8]['zero_point'])
+            out = torch.ops.quantized.add(outputs, perturbed_outputs, self.quant_dict[9]['scale'], self.quant_dict[9]['zero_point'])
+            out = out.dequantize()
 
-        # perturbed feedforward
-        x = torch.ops.quantized.mul(x, sign_input, default_scale, default_zero_point)
-
-        perturbed_outputs = torch.nn.quantized.functional.conv2d(x,
-                            weight=delta_kernel, bias=bias, stride=self.stride, padding=self.padding,
-                            dilation=self.dilation, groups=self.groups, scale=default_scale, zero_point=default_zero_point)
-        perturbed_outputs = torch.ops.quantized.mul(perturbed_outputs, sign_output, default_scale, default_zero_point)
-        out = torch.ops.quantized.add(outputs, perturbed_outputs, default_scale, default_zero_point)
+        else:
+            if x.dtype!=torch.quint8:
+                x = torch.quantize_per_tensor(x, default_scale, default_zero_point, torch.quint8)
+
+            outputs = torch.nn.quantized.functional.conv2d(x, self.quantized_mu_weight, bias, self.stride, self.padding,
+                            self.dilation, self.groups, scale=default_scale, zero_point=default_zero_point) # input: quint8, weight: qint8, bias: fp32
+
+            # sampling perturbation signs
+            sign_input = torch.zeros(x.shape).uniform_(-1, 1).sign()
+            sign_output = torch.zeros(outputs.shape).uniform_(-1, 1).sign()
+            sign_input = torch.quantize_per_tensor(sign_input, default_scale, default_zero_point, torch.quint8)
+            sign_output = torch.quantize_per_tensor(sign_output, default_scale, default_zero_point, torch.quint8)
+
+            # getting perturbation weights
+            eps_kernel = torch.quantize_per_tensor(self.eps_kernel.data.normal_(), normal_scale, 0, torch.qint8)
+            new_scale = (self.quantized_sigma_weight.q_scale())*(eps_kernel.q_scale())
+            delta_kernel = torch.ops.quantized.mul(self.quantized_sigma_weight, eps_kernel, new_scale, 0)
+
+            bias = None
+            if self.bias:
+                eps_bias = self.eps_bias.data.normal_()
+                bias = (self.quantized_sigma_bias * eps_bias)
+
+            # perturbed feedforward
+            x = torch.ops.quantized.mul(x, sign_input, default_scale, default_zero_point)
+
+            perturbed_outputs = torch.nn.quantized.functional.conv2d(x,
+                                weight=delta_kernel, bias=bias, stride=self.stride, padding=self.padding,
+                                dilation=self.dilation, groups=self.groups, scale=default_scale, zero_point=default_zero_point)
+            perturbed_outputs = torch.ops.quantized.mul(perturbed_outputs, sign_output, default_scale, default_zero_point)
+            out = torch.ops.quantized.add(outputs, perturbed_outputs, default_scale, default_zero_point)
 
         if return_kl:
             return out, 0