量化算子

稳定 API

torch.ops.fbgemm.FloatOrHalfToFusedNBitRowwiseQuantizedSBHalf(input, bit_rate) → Tensor

使用行级量化将 FP32/16 转换为 INT8/4/2。

参数：

• input (Tensor) – 输入张量。必须是 FP32（torch.float）或 FP16（torch.half），并且必须是二维。

• bit_rate (int) – 量化比特率（INT2 为 2，INT4 为 4，或 INT8 为 8）

返回值：

量化输出（张量）。数据类型为 torch.uint8（字节类型）

示例

>>> # Randomize input
>>> input = torch.randn(2, 4, dtype=torch.float32, device="cuda")
>>> print(input)
tensor([[ 0.8247,  0.0031, -1.0068, -1.2081],
        [ 0.5427,  1.5772,  1.0291, -0.7626]], device='cuda:0')
>>> # Quantize
>>> output = torch.ops.fbgemm.FloatOrHalfToFusedNBitRowwiseQuantizedSBHalf(input, bit_rate=4)
>>> print(output)
tensor([[159,   1,  86,  48, 213, 188],
        [248,  11, 254,  48,  26, 186]], device='cuda:0', dtype=torch.uint8)

其他 API