[docs]classBatchNorm2d(_BatchNorm):r"""This is the quantized version of :class:`~torch.nn.BatchNorm2d`."""_NNI_BN_RELU_MODULE=nni.BNReLU2ddef__init__(self,num_features,eps=1e-5,momentum=0.1,device=None,dtype=None)->None:factory_kwargs={"device":device,"dtype":dtype}super().__init__(num_features,eps,momentum,**factory_kwargs)def_get_name(self):return"QuantizedBatchNorm2d"def_check_input_dim(self,input):# Temporarily using len(shape) instead of ndim due to JIT issue# https://github.com/pytorch/pytorch/issues/23890iflen(input.shape)!=4:raiseValueError("Input shape must be `(N, C, H, W)`!")defforward(self,input:torch.Tensor)->torch.Tensor:# disabling this since this is not symbolically traceable# self._check_input_dim(input)returntorch.ops.quantized.batch_norm2d(input,self.weight,self.bias,self.running_mean,self.running_var,self.eps,self.scale,self.zero_point,)@classmethoddeffrom_float(cls,mod,use_precomputed_fake_quant=False):return_BatchNorm.from_float(cls,mod,use_precomputed_fake_quant=use_precomputed_fake_quant)
[docs]classBatchNorm3d(_BatchNorm):r"""This is the quantized version of :class:`~torch.nn.BatchNorm3d`."""_NNI_BN_RELU_MODULE=nni.BNReLU3ddef__init__(self,num_features,eps=1e-5,momentum=0.1,device=None,dtype=None):factory_kwargs={"device":device,"dtype":dtype}super().__init__(num_features,eps,momentum,**factory_kwargs)def_get_name(self):return"QuantizedBatchNorm3d"def_check_input_dim(self,input):# Temporarily using len(shape) instead of ndim due to JIT issue# https://github.com/pytorch/pytorch/issues/23890iflen(input.shape)!=5:raiseValueError("Input shape must be `(N, C, H, W)`!")defforward(self,input:torch.Tensor)->torch.Tensor:# disabling this since this is not symbolically traceable# self._check_input_dim(input)returntorch.ops.quantized.batch_norm3d(input,self.weight,self.bias,self.running_mean,self.running_var,self.eps,self.scale,self.zero_point,)@classmethoddeffrom_float(cls,mod,use_precomputed_fake_quant=False):return_BatchNorm.from_float(cls,mod,use_precomputed_fake_quant=use_precomputed_fake_quant)
Docs
Access comprehensive developer documentation for PyTorch
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.