PyTorch torch.fake_quantize_per_channel_affine function


Pytorch torch 参考手册Pytorch torch Reference Manual

torch.fake_quantize_per_channel_affineIt is a function in PyTorch used for channel-wise fake quantization of tensors, commonly used in quantization-aware training.

Function definition

torch.fake_quantize_per_channel_affine(input, scale, zero_point, axis, quant_min, quant_max)

Usage example

Example

import torch

# Create input tensor
x = torch.randn(2, 3, 4, 5)

# Define scale and zero point
scale = torch.tensor([1.0, 1.2, 1.5])
zero_point = torch.tensor([0, 0, 0], dtype=torch.long)

# Perform channel-wise fake quantization
y = torch.fake_quantize_per_channel_affine(x, scale, zero_point, axis=1, quant_min=0, quant_max=255)
print("Quantized shape:", y.shape)

Pytorch torch 参考手册Pytorch torch Reference Manual

Other extensions