quantize_affine

torchao.quantization.quantize_affine(input: Tensor, block_size: Tuple[int, ...], scale: Tensor, zero_point: Optional[Tensor], output_dtype: dtype, quant_min: Optional[Union[int, float]] = None, quant_max: Optional[Union[int, float]] = None, zero_point_domain: Optional[ZeroPointDomain] = ZeroPointDomain.INT) → Tensor

参数:

• input (torch.Tensor) – 原始 float32、float16 或 bfloat16 张量

• block_size – (Tuple[int, …]): 量化粒度，这意味着共享相同 qparam 的张量元素的大小。例如，当大小与输入张量维度相同时，我们使用per tensor quantization（每张量量化）

• scale (float) – 仿射量化的量化参数

• zero_point (int) – 仿射量化的量化参数

• output_dtype (torch.dtype) – 输出张量请求的数据类型 (例如 torch.uint8)

• quant_min (Optional[int]) – 输出张量的最小量化值，如果未指定，则将从 dtype 派生

• quant_max (Optional[int]) – 输出张量的最大量化值，如果未指定，则将从 dtype 派生

• zero_point_domain (ZeroPointDomain) – zero_point 所在的域，如果 zero_point 在整数域中，则应为整数或浮点数，量化期间将 zero point 添加到量化的整数值；如果 zero_point 在浮点域中，则量化期间将 zero point 从浮点（未量化）值中减去。默认值为 ZeroPointDomain.INT

注意

block_size 如何表示不同的粒度？假设我们有一个大小为 (3, 3, 10, 10) 的张量，下表显示了 block_size 如何表示不同的粒度

粒度类型 | block_size

per_tensor（每张量） | (3, 3, 10, 10) per_axis (axis=0)（每轴（轴=0）） | (1, 3, 10, 10) per_axis (axis=1)（每轴（轴=1）） | (3, 1, 10, 10)

per_group (groupsize=2)（每组 (groupsize=2)） | (3, 3, 10, 2) per_group (groupsize=2) for axis = 3（轴 = 3 的每组 (groupsize=2)） | (3, 3, 2, 10)

输出

具有请求数据类型的量化张量