4.2.4.3. Supported Operators

4.2.4.3.1. Supported Torch Operator

The operators in the table below have the following restrictions unless otherwise specified: input/output data type int8, input_shape: [N, C, H, W], input_size <1 G bytes, 1<=N<=4096, 1<=H, W, C<=65536, Feature dimension: 4. If additional restriction information is added to the table, the table shall prevail.

Torch ops

To prepare the floating point model, replace it with

Bernoulli2 Support

Bayes/Bayes-e Support

torch.add

torch.nn.quantized.FloatFunctional or horizon.nn.quantized.FloatFunctional

Supported. QAT has training parameters and cannot be used alone in prediction.
in_channel<=2048

Support. QAT has training parameters and cannot be used alone in prediction.
Input and output support int8/int16. Supports broadcasts other than N dimensions, and there can only be one input broadcast.

torch.sub

horizon.nn.quantized.FloatFunctional

Supported. QAT has training parameters and cannot be used alone in prediction.
in_channel<=2048

Support. QAT has training parameters and cannot be used alone in prediction.
Input and output support int8/int16. Supports broadcasts other than N dimensions, and there can only be one input broadcast.

torch.mul

torch.nn.quantized.FloatFunctional or horizon.nn.quantized.FloatFunctional

Supported. QAT has training parameters and cannot be used alone in prediction.
in_channel<=2048

Support. QAT has training parameters and cannot be used alone in prediction.
Input and output support int8/int16. Supports broadcasts other than N dimensions, and there can only be one input broadcast.

torch.sum

horizon.nn.quantized.FloatFunctional

only supports sum in batch and channel directions. QAT has
training parameters, do not use it in prediction alone.

Support. QAT has training parameters and cannot be used alone in prediction.
Input and output support int8/int16. Only support HWC three-dimensional sum

torch.matmul

horizon.nn.quantized.FloatFunctional

Supported. For matmul(a, b),
b can be transposed internally by setting the argument

support. Input int8, output int8/int16/int32.
For matmul(a, b), you can transpose
b internally by setting parameters. input shape: [N, C, H, W],
input_size<1 G bytes, N<=4096, C, H, W<=8192.

torch.cat

torch.nn.quantized.FloatFunctional or horizon.nn.quantized.FloatFunctional

For cat(a, b), the difference between a and b should not be too large, otherwise
one of the operands will eat the other An operand phenomenon.
QAT has training parameters, do not use alone in prediction.

For cat(a, b), the difference between a and b should not be too large, otherwise
one operand will eat the other operand.
QAT has training parameters, do not use alone in prediction. input shape: [N, C, H, W],
N<=4096, HWC<=65536, 2<=input number<=1024

torch.maximum

horizon.nn.quantized.FloatFunctional

Not supported

Support input and output int8/int16

torch.minimum

horizon.nn.quantized.FloatFunctional

Not supported

Support input and output int8/int16

torch.mean

horizon.nn.quantized.FloatFunctional

Only the mean in the channel direction is supported. The training parameters are provided in the QAT and cannot be used alone in prediction.

Find mean on the CHW.The training parameters are provided in the QAT and cannot be used alone in prediction. The input/output supports int8/int16.

torch.sqrt

horizon.nn.Sqrt

Not supported

Supports input and output int8/int16. Low-level look-up table implementation, there is a risk of accuracy

torch.atan

horizon.nn.Atan

Not supported

Supports input and output int8/int16. Low-level look-up table implementation, there is a risk of accuracy

torch.sin

horizon.nn.Sin

Not supported

Supports input and output int8/int16. Low-level look-up table implementation, there is a risk of accuracy

torch.cos

horizon.nn.Cos

Not supported

Supports input and output int8/int16. Low-level look-up table implementation, there is a risk of accuracy

torch.clamp/clip, tensor.clamp/clip

Not supported

Supported inputs for min and max are Tensor/constant Tensor/scalar/None. When it is a constant Tensor, the input data range of min and max should be consistent with the input, otherwise there is a risk of precision

torch.pow

horizon.nn.Pow

Not supported

Supports input and output int8/int16. Low-level look-up table implementation, there is a risk of accuracy

torch.max

support. Only available as model output. The output format is different from torch: the output supported by the compiler
is a Tensor,
the value in one channel is max_value,
the value in the other channel is max_value_index

supported. Input supports int8/int16. One output is int8/16 value and the other is int32 index. The index can only be available as model output. Input_shape: [N, C, H, W], 1<=N<=4096, 1<=H, W, C<=65535

torch.min

Not supported

Refer to torch.max

tensor.max

Refer to torch.max

Refer to torch.max

tensor.min

Not supported

Refer to torch.min

torch.split

Supported

Input/output int8/int16

torch.eq

Not supported

Input/output int8/int16

torch.ge

Not supported

Input/output int8/int16

torch.greater

Not supported

Input/output int8/int16

torch.greater_equal

Not supported

Input/output int8/int16

torch.gt

Not supported

Input/output int8/int16

torch.le

Not supported

Input/output int8/int16

torch.less

Not supported

Input/output int8/int16

torch.less_equal

Not supported

Input/output int8/int16

torch.lt

Not supported

Input/output int8/int16

tensor.eq

Not supported

Input/output int8/int16

tensor.ge

Not supported

Input/output int8/int16

tensor.greater

Not supported

Input/output int8/int16

tensor.greater_equal

Not supported

Input/output int8/int16

tensor.gt

Not supported

Input/output int8/int16

tensor.le

Not supported

Input/output int8/int16

tensor.less

Not supported

Input/output int8/int16

tensor.less_equal

Not supported

Input/output int8/int16

tensor.expand

Not supported

Input/output int8/int16

tensor.repeat

Not supported

Input/output int8/int16

torch.nn.GLU

Not supported

Support input and output int8/int16

tensor.tile

Not supported

Input/output int8/int16

torch.nn.GELU

Not supported

Input/output int8/int16

torch.nn.LSTMCell

Not supported

Input/output int8/int16.input is 2 dimension

torch.nn.LeakyReLU

Not supported

Supports input and output int8/int16. Low-level look-up table implementation, there is a risk of accuracy

torch.nn.PReLU

Not supported

Input/output int8/int16

torch.argmax

Refer to torch.max

Refer to torch.max

torch.argmin

Refer to torch.min

Refer to torch.min

tensor.argmax

Refer to torch.max

Refer to torch.max

tensor.argmin

Refer to torch.min

Refer to torch.min

tensor.reshape

Only the reshape in H and W direction can be supported.

dim <= 10, 1 <= each_dim_size <= 65536

torch.nn.Tanh

Supported. Underlying implementation of lookup tables with some accuracy risk.

Supports input and output int8/int16

torch.nn.ReLU

Conv2d+BN+ReLU This mode will automatically fuse into
BpuConv2d, otherwise it’s not friendly for QAT to to run ReLU alone.
It is recommended to use ReLU6. During the QAT,
use relu6 by default.

Conv2d+BN+ReLU This mode will automatically fuse into BpuConv2d, otherwise it’s not friendly for QAT to to run ReLU alone.
It is recommended to use ReLU6. During the QAT,
use relu6 by default.

torch.nn.SiLU

Supported. Underlying implementation of lookup tables.
Some accuracy risk

Supported

torch.nn.Conv2d

Supported. kernel<=7.channel(one group) <= 2048.
dilation=(1, 1)/(2, 2)/(4, 4), when dilation!=(1, 1),
stride must be (1, 1). HWC <= 32768

Supported. out_channel<=8192, when output as a model, out_channel <= 16384.
input channel<=8192, kernel<32, dilation<=16, when dilation!=1, stride must be
=1. It supports sumin, and conv with sumin only supports stride as (1, 1) or (2, 2).
weight_shape: [N, C, H, W], N, C<=8192, H, W<=31, as model output C<=16384,
weight_size < 65535.padding<=256.

torch.nn.Linear

Not supported

Supported. in_features <= 8192, out_features <= 8192.

torch.nn.Conv3d

Not supported

input: [N, C, D, H, W] int8, N<=128; H, W, D, C<=65536;
weight: [C_o, C_i, D, H, W] int8, N, C<=65536, D, H<=9, W<=8191;
bias: int32;
output: [N, C, D, H, W] int8, int16, int32;
stride: [D, H, W], D, H, W in (1, 2), and D == H == W;
padding: [D, H, W], D<=kernel_d/2, H<=kernel_h/2, W<=kernel_w/2(indicates weight W size)
group, dilation: not supported yet

tensor.transpose

Not supported

Input/output int8, int16, int32.Transpose for N dimensions is not supported.

torch.nn.Sigmoid

Supported. Underlying implementation of lookup tables.
Some accuracy risk.

Supports input and output int8/int16

torch.nn.Dropout

Training op, not reflected in the prediction model.

Training op, not reflected in the prediction model.

torch.nn.Softmax

Not supported

Input/output int8/int16. Underlying implementation of lookup tables.
Some accuracy risk.

torch.nn.Identity

Training op, not reflected in the prediction model.

Training op, not reflected in the prediction model.

torch.nn.AvgPool2d

1<=kernel<=7, 1<=stride<=185

Supported. 1<=kernel, stride, padding<=256;

torch.nn.MaxPool2d

1<=kernel<=64, 1<=stride<=256, padding>=0

input_shape: [N, C, H, W], 1<=H, W, C<=8192;1<=kernel, stride<=256;
0<=padding<=255;

torch.nn.ZeroPad2d

Supported.

Input/output int8/int16.

torch.nn.Dropout2d

Training op, not reflected in the prediction model.

Training op, not reflected in the prediction model.

torch.nn.Layernorm

Not supported

Input/output int8/int16, Underlying implementation of lookup tables.
User can modify the underlying rsqrt LUT param by rsqrt_kwargs attribute.
Please try layernorm_op.rsqrt_kwargs = {"auto_divide_strategy": "curvature"} if model accuracy reduced after convert.
Some accuracy risk.
H * W <= 16384, normalized_shape H * W < 16384

torch.nn.BatchNorm2d

BatchNorm2d. Absorbed at the QAT stage and not reflected in the
prediction model. Due to compiler limitations, the
BatchNorm2d used independently calls the BpuConvolution implementation from the bottom.

BatchNorm2d. Absorbed at the QAT stage and
not reflected in the model.
Standalone usage limiations refer to Conv2d.

torch.nn.BatchNorm3d

BatchNorm3d. Absorbed at the QAT stage and not reflected in the
prediction model. Due to compiler limitations, the
BatchNorm3d used independently calls the BpuConvolution implementation from the bottom.

BatchNorm3d. Absorbed at the QAT stage and
not reflected in the model.
Standalone usage limiations refer to Conv2d.

Tensor.__getitem__

Supported

Supported

torch.nn.ConstantPad2d

Supported

Supported. Limitations refer to ZeroPad2d.

torch.nn.SyncBatchNorm

Training op, not reflected in the prediction model.

Training op, not reflected in the prediction model.

torch.nn.ChannelShuffle

Supported

Input/output int8/int16. The values in shuffle_index cannot be duplicated.

torch.nn.PixelShuffle

Supported. Input/output int8/int16.

Supported. Input/output int8/int16.

torch.nn.PixelUnshuffle

Supported. Input/output int8/int16.

Supported. Input/output int8/int16.

torch.nn.ConvTranspose2d

Supported. 2<=kernel<= 14.channel<=2048.
padding H*W=[0, (kernel_h-1)/2] * [0, (kernel_w-1)/2]
2<=stride<=4, dilation=(1, 1)

Supported. Input shape: [N, C, H, W], 1<=N<=128, 1<=channel<=2048;
weight_shape: [N, C, H, W], 1<=N, C<=2048, 2<=H, W<=14,
weight_size<=65535;kernel>=stride, 1<=stride<=14, 1<=out_channel<=2048,
in_channel<=2048
pad<=kernel/stride, 0<=out_pad<=1;bias type: int32; It supports sumin,
sumin input type: int8;0<=output_padding<=1; It supports group. Both weight_n and
input channels are required to be divisible by group;dilation=1

torch.nn.Upsample

Supported. Refer to torch.nn.functional.interpolate

Refer to torch.nn.functional.interpolate

torch.nn.UpsamplingNearest2d

Supported. The limitations refer to torch.nn.functional.interpolate

Supported. The limitations refer to torch.nn.functional.interpolate

torch.nn.UpsamplingBilinear2d

Supported. The limitations refer to torch.nn.functional.interpolate

Supported. The limitations refer to torch.nn.functional.interpolate

torch.nn.functional.pad

Support pad mode except reflect

Support pad mode except reflect

torch.nn.functional.relu

torch.nn.ReLU

Conv2d+BN+ReLU This mode will automatically fuse into
BpuConv2d, otherwise it’s not friendly for QAT
to run ReLU.ReLU alone. It is recommended to use ReLU6. During the QAT,
use relu6 by default.

Conv2d+BN+ReLU This mode will automatically fuse into
BpuConv2d, otherwise it’s not friendly for QAT
to run ReLU.ReLU alone. It is recommended to use ReLU6. During the QAT,
use relu6 by default.

torch.nn.functional.relu6(fused)

torch.nn.ReLU6

Conv2d+BN+ReLU6 This mode will automatically fuse into BpuConv2d,
otherwise run ReLU6 alone.

Conv2d+BN+ReLU6 This mode will automatically fuse into BpuConv2d,
otherwise run ReLU6 alone.

torch.nn.ReplicationPad2d

Supported

Supported. The limitations refer to ZeroPad2d.

torch.quantization.QuantStub

horizon.quantization.QuantStub

Supported. Typical usage scenario: Network model input.
Model segmentation scenario: The data needs to be quantized before it is sent from the CPU to the BPU
. Scale parameter settings:
The scale settings are related to the specific input.
The settings are used to quantify the input data of the float type to the int8 type with high accuracy as much as possible. There are two requirements:
It can cover all (at least most) input data ,
High quantization accuracy. For example: the input float range is (-1, 1).
Then, let’s set scale = 1 / 128. Float pre-training model:
In the pre-training model, since the model has been trained,
it’s not necessarily follow the above the parameter setting method of scale.
In this case, it can be solved by inserting a special conv.
It requires that the data distribution in QuantStub should be uniform

Supported.Typical usage scenario: Network model input.
Model segmentation scenario: The data needs to be quantized before it is sent from the CPU to the BPU
. Scale parameter settings:
The scale settings are related to the specific input.
The settings are used to quantify the input data of the float type to the int8 type with high accuracy as much as possible. There are two requirements:
It can cover all (at least most) input data ,
High quantization accuracy. For example: the input float range is (-1, 1).
Then, let’s set scale = 1 / 128. Float pre-training model:
In the pre-training model, since the model has been trained,
it’s not necessarily follow the above the parameter setting method of scale.
In this case, it can be solved by inserting a special conv.
It requires that the data distribution in QuantStub should be uniform.

torch.quantization.DeQuantStub

Typical usage scenario: In the scenario of network model segmentation, it is necessary to transfer data
from BPU to CPU for inverse quantization on the CPU,
to facilitate processing on the CPU.

Typical usage scenario: In the scenario of network model segmentation, it is necessary to transfer data
from BPU to CPU for inverse quantization on the CPU,
to facilitate processing on the CPU.

torch.nn.functional.interpolate

horizon.nn.Interpolate

Only nearest and billinear interpolation modes are supported. 1/256

Only nearest and billinear interpolation modes are supported. Input_shape: [N, C, H, W], 1<=C, H, W<=8192

torch.nn.functional.grid_sample

Not supported

Supported. Input shape: [N, C, H, W], 1<=H, W<=1024 and HW<=5121024;
grid only supports qint16. It only supports bilinear and nearest interpolation
The padding mode only supports zeros and border;

No Correlation operator in torch

horizon.nn.Correlation

Not supported

Supported. Input shape: [N, C, H, W], 1<=C<=1024, kernel must be odd number,
min(H, W) + 2 * pad - 2 * (max_displacement/stride2*stride2)-kernel+1 >= stride1

torch.log

horizon.nn.HardLog

Not supported

Supports input and output int8/int16. Low-level look-up table implementation, there is a risk of accuracy

torch.masked_fill

Not supported

Supports input and output int8/int16

No PointPillarsScatter operator in torch

horizon.nn.PointPillarsScatter

Not supported

Supports input and output int8/int16. Requires three inputs, the shapes of the three inputs are: [M, C], [M, 4] and [4], and the third input represents the shape of the output feature of the operator: (N, C, H, W)

torch.div

horizon.nn.Div

Not supported

Supports input and output int16

No bgr2centered_yuv operator in torch

horizon.bgr2centered_yuv

Not supported

Supported. Used for data preprocessing. Convert BGR images of 0~255 to centered YUV

No rgb2bgr operator in torch

horizon.rgb2bgr

Not supported

Supported. Used for data preprocessing. Convert RGB images of shape [N, 3, H, W] to BGR

No bgr2rgb operator in torch

horizon.bgr2rgb

Not supported

Supported. Used for data preprocessing. Convert BGR images of shape [N, 3, H, W] to RGB

No bgr2yuv operator in torch

horizon.bgr2yuv

Not supported

Supported. Used for data preprocessing. Convert BGR images of 0~255 to YUV

No rgb2yuv operator in torch

horizon.rgb2yuv

Not supported

Supported. Used for data preprocessing. Convert RGB images of 0~255 to YUV

No bgr2centered_yuv operator in torch

horizon.bgr2centered_yuv

Not supported

Supported. Used for data preprocessing. Convert BGR images of 0~255 to centered YUV

No rgb2centered_yuv operator in torch

horizon.rgb2centered_yuv

Not supported

Supported. Used for data preprocessing. Convert RGB images of 0~255 to centered YUV

No bgr2gray operator in torch

horizon.bgr2gray

Not supported

Supported. Used for data preprocessing. Convert BGR images of 0~255 to gray

No rgb2gray operator in torch

horizon.rgb2gray

Not supported

Supported. Used for data preprocessing. Convert RGB images of 0~255 to gray

No bgr2centered_gray operator in torch

horizon.bgr2centered_gray

Not supported

Supported. Used for data preprocessing. Convert BGR images of 0~255 to centered gray

No rgb2centered_gray operator in torch

horizon.rgb2centered_gray

Not supported

Supported. Used for data preprocessing. Convert RGB images of 0~255 to centered gray

No centered_yuv2bgr operator in torch

horizon.centered_yuv2bgr

Not supported

Supports input and output int8. Used when users have trained QAT model with BGR images. Inserts this operator in the quantized model to convert centered YUV inputs to BGR

No centered_yuv2rgb operator in torch

horizon.centered_yuv2rgb

Not supported

Supports input and output int8. Used when users have trained QAT model with RGB images. Inserts this operator in the quantized model to convert centered YUV inputs to RGB

torch.nn.MultiheadAttention

Not supported

Do not support add_bias_kv, add_zero_attn and qkv with different embed_dim. Supports input and output int8/int16. with low_level look-up table implementation and mask quantization, there is a risk of accuracy.

torch.reciprocal

horizon.nn.Reciprocal

Not supported

Supports input and output int8/int16

No rcnn_post_process operator in torch

horizon.nn.RcnnPostProcess

Not supported

This operator can only be deployed on CPU when inference. Supports input and output float. Includes a post-process pipeline such as NMS for the output of RCNN.

torch.topk

Not supported

Supports input and output int8/int16/int32.

torch.gather

Not supported

Supports input and output int8/int16/int32.

torch.abs

Not supported

Supports input and output int8/int16.

torch.nn.Softplus

Not supported

Supports input and output int8/int16. Low-level look-up table implementation, there is a risk of accuracy.

torch.nn.ELU

Not supported

Supports input and output int8/int16. Low-level look-up table implementation, there is a risk of accuracy.

torch.ceil

horizon.nn.Ceil

Not supported

Supports input and output int8/int16. int8 supports input magnitude not greater than 1e6, while int16 supports input magnitude not greater than 1e8.

torch.floor

horizon.nn.Floor

Not supported

Supports input and output int8/int16. int8 supports input magnitude not greater than 1e6, while int16 supports input magnitude not greater than 1e8.

no MultiScaleDeformableAttention operator in torch

horizon.nn.MultiScaleDeformableAttention

Not supported

Supports input and output int8.

torch.nn.HardSigmoid

Not supported

Supports input and output int8/int16. Low-level look-up table implementation, there is a risk of accuracy.

torch.acos

horizon.nn.Acos

Not supported

Supports input and output int8/int16. Low-level look-up table implementation, there is a risk of accuracy

torch.acosh

horizon.nn.Acosh

Not supported

Supports input and output int8/int16. Low-level look-up table implementation, there is a risk of accuracy

torch.asin

horizon.nn.Asin

Not supported

Supports input and output int8/int16. Low-level look-up table implementation, there is a risk of accuracy

torch.asinh

horizon.nn.Asinh

Not supported

Supports input and output int8/int16. Low-level look-up table implementation, there is a risk of accuracy

torch.atanh

horizon.nn.Atanh

Not supported

Supports input and output int8/int16. Low-level look-up table implementation, there is a risk of accuracy

torch.cosh

horizon.nn.Cosh

Not supported

Supports input and output int8/int16. Low-level look-up table implementation, there is a risk of accuracy

torch.erf

horizon.nn.Erf

Not supported

Supports input and output int8/int16. Low-level look-up table implementation, there is a risk of accuracy

torch.selu

horizon.nn.Selu

Not supported

Supports input and output int8/int16. Low-level look-up table implementation, there is a risk of accuracy

torch.sinh

horizon.nn.Sinh

Not supported

Supports input and output int8/int16. Low-level look-up table implementation, there is a risk of accuracy

torch.tan

horizon.nn.Tan

Not supported

Supports input and output int8/int16. Low-level look-up table implementation, there is a risk of accuracy

4.2.4.3.2. Supported Torchvision Operator

torchvision Operator

To prepare the floating point model, replace it with

Bernoulli2 Support

Bayes/Bayes-e Support

torchvision.ops.RoIAlign

Supported

Supported. 1<=feature number<=5;bbox need to be List[Tensor], shape:[1, box_num, 4],
The 6 numbers in the last dimension of bbox are: [left, top, right, bottom]

torchvision.ops.MultiScaleRoIAlign

horizon.nn.MultiScaleRoIAlign

Supported

Supported. The limitations refer to RoIAlign.

torchvision.models.detection.rpn.AnchorGenerator

horizon.nn.AnchorGenerator

Only the case where Tensor.shape can be determined offline is supported.

Input: int8/int16/int32/float32, output: float32