fake_quantize_affine¶
- torchao.quantization.fake_quantize_affine(input: Tensor, block_size: Tuple[int, ...], scale: Tensor, zero_point: Optional[Tensor], quant_dtype: dtype, quant_min: Optional[Union[int, float]] = None, quant_max: Optional[Union[int, float]] = None, zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT) Tensor[source]¶
General fake quantize op for quantization-aware training (QAT). This is equivalent to calling quantize_affine + dequantize_affine but without the dtype casts.
- Parameters:
input (torch.Tensor) – original float32, float16 or bfloat16 Tensor
block_size – (Tuple[int, …]): granularity of quantization, this means the size of the tensor elements that’s sharing the same qparam e.g. when size is the same as the input tensor dimension, we are using per tensor quantization
scale (float) – quantization parameter for affine quantization
zero_point (int) – quantization parameter for affine quantization
quant_dtype (torch.dtype) – desired quantized dtype for determining and validating quant_min and quant_max values.
quant_min (Optional[int]) – minimum quantized value for output Tensor, if not specified, it will be derived from dtype
quant_max (Optional[int]) – maximum quantized value for output Tensor, if not specified, it will be derived from dtype
zero_point_domain (ZeroPointDomain) – the domain that zero_point is in, should be either integer or float if zero_point is in integer domain, zero point is added to the quantized integer value during quantization if zero_point is in floating point domain, zero point is subtracted from the floating point (unquantized) value during quantization default is ZeroPointDomain.INT
