[docs]classLinear(nn.Linear):r""" A linear module attached with FakeQuantize modules for weight, used for quantization aware training. We adopt the same interface as `torch.nn.Linear`, please see https://pytorch.org/docs/stable/nn.html#torch.nn.Linear for documentation. Similar to `torch.nn.Linear`, with FakeQuantize modules initialized to default. Attributes: weight: fake quant module for weight """_FLOAT_MODULE=nn.Lineardef__init__(self,in_features,out_features,bias=True,qconfig=None,device=None,dtype=None,)->None:factory_kwargs={"device":device,"dtype":dtype}super().__init__(in_features,out_features,bias,**factory_kwargs)assertqconfig,"qconfig must be provided for QAT module"self.qconfig=qconfigself.weight_fake_quant=qconfig.weight(factory_kwargs=factory_kwargs)defforward(self,input):returnF.linear(input,self.weight_fake_quant(self.weight),self.bias)
[docs]@classmethoddeffrom_float(cls,mod,use_precomputed_fake_quant=False):r"""Create a qat module from a float module or qparams_dict Args: `mod` a float module, either produced by torch.ao.quantization utilities or directly from user """asserttype_before_parametrizations(mod)==cls._FLOAT_MODULE,(" qat."+cls.__name__+".from_float only works for "+cls._FLOAT_MODULE.__name__)asserthasattr(mod,"qconfig"),"Input float module must have qconfig defined"assertmod.qconfig,"Input float module must have a valid qconfig"iftype_before_parametrizations(mod)==LinearReLU:mod=mod[0]qconfig=mod.qconfigqat_linear=cls(mod.in_features,mod.out_features,bias=mod.biasisnotNone,qconfig=qconfig,)ifis_parametrized(mod,"weight"):transfer_parametrizations_and_params(mod,qat_linear,"weight")else:qat_linear.weight=mod.weightifis_parametrized(mod,"bias"):transfer_parametrizations_and_params(mod,qat_linear,"bias")else:qat_linear.bias=mod.biasreturnqat_linear
To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.