File size: 2,565 Bytes
96ba80e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | import torch
import torch.nn as nn
import torch.nn.functional as F
from quantize.quantizer import UniformAffineQuantizer
class QuantLinear(nn.Module):
"""
Quantized Module that can perform quantized convolution or normal convolution.
To activate quantization, please use set_quant_state function.
"""
def __init__(
self,
org_module: nn.Linear,
weight_quant_params: dict = {},
act_quant_params: dict = {},
disable_input_quant=False,
):
super().__init__()
self.fwd_kwargs = dict()
self.fwd_func = F.linear
self.register_buffer('weight',org_module.weight)
if hasattr(org_module,"bias") and org_module.bias is not None:
self.register_buffer('bias',org_module.bias)
else:
self.bias = None
if hasattr(org_module,"in_features"):
self.in_features = org_module.in_features
self.out_features = org_module.out_features
# de-activate the quantized forward default
self.use_weight_quant = False
self.use_act_quant = False
self.quant_rate = 1.0
# import ipdb;ipdb.set_trace()
# initialize quantizer
if weight_quant_params["n_bits"] > 1:
self.weight_quantizer = UniformAffineQuantizer(**weight_quant_params,shape=org_module.weight.shape,is_weight_quant=True)
if not disable_input_quant:
self.act_quantizer = UniformAffineQuantizer(**act_quant_params)
else:
self.act_quantizer = None
self.disable_input_quant = disable_input_quant
self.use_temporary_parameter = False
def forward(self, input: torch.Tensor):
if self.use_temporary_parameter:
weight = self.weight_quantizer(self.temp_weight,self.quant_rate)
bias = self.temp_bias
elif self.use_weight_quant:
weight = self.weight_quantizer(self.weight,self.quant_rate)
bias = self.bias
else:
weight = self.weight
bias = self.bias
if self.use_act_quant and not self.disable_input_quant:
input = self.act_quantizer(input,self.quant_rate)
out = self.fwd_func(input, weight, bias, **self.fwd_kwargs)
return out
def set_quant_state(self, weight_quant: bool = False, act_quant: bool = False, quant_rate:float = 1.0):
self.use_weight_quant = weight_quant
self.use_act_quant = act_quant
self.quant_rate = quant_rate
|