精度调优｜conv+depth2space 替换 resize 指导

1.技术背景

在进行模型压缩与加速的过程中，量化技术成为了提升推理速度和降低计算资源消耗的重要手段。然而，在实际应用中，许多用户发现，采用 resize 操作时，仅能使用 int8 精度量化，这一限制导致了模型精度的显著下降。尽管 int8 精度在提升计算效率方面具有优势，但精度的丧失却使得模型的推理结果偏差增大，给实际应用带来了不少困扰。如何在保证性能的同时，最大程度地减少精度损失，成为了当前技术实现中的一个难题。 在当前工具链版本下（J6 3.0.31），resize 算子仅支持 int8 量化精度，不支持 int16，因此该类算子有一定概率触发精度下降问题。

针对这一问题，本文将介绍一种新的方法，可以有效提升上采样操作中的精度，解决传统 int8 精度量化带来的精度下降问题。通过巧妙地优化模型结构，能够在不显著影响计算效率的前提下，显著提高上采样的精度。 本文介绍的这种解决方案，在不影响模型权重（无需重训）的情况下，通过算子替换，使得上采样功能支持 int16 量化精度，以解决精度下降问题。 相信这一方案将为广大开发者带来帮助～

2.方案介绍

onnx 中的 resize 算子，在 pytorch 代码中常表现为 F.interpolate 函数。当 F.interpolate 的 mode 为 nearest 时，该函数的功能和 conv+depth2space 完全等效。而 conv 和 depth2space 均支持 int16 量化，因此可以通过算子替换的方式变相实现上采样的 int16 支持。

import torch
import torch.nn as nn
import torch.nn.functional as F
class Conv2DInterpolate(nn.Module):
    def 
__init__
(self, inputs_channel=1, scale_factor=2) -> None:
        super().
__init__
()
        self.conv = nn.Conv2d(
            in_channels=inputs_channel,
            out_channels=inputs_channel * (scale_factor**2),
            kernel_size=3,
            bias=False,
            padding=1,
        )
        self.scale_factor = scale_factor
        self.inputs_channel = inputs_channel
        self.depth2space = torch.nn.PixelShuffle(scale_factor)
        self._init_weights()
    def _init_weights(self):
        conv_weight = torch.zeros(
            self.conv.weight.size(),
            dtype=self.conv.weight.dtype,
        )
        num_conv = conv_weight.shape[0]
        for i_N in range(num_conv):
            i_c = i_N // (self.scale_factor**2)
            conv_weight[i_N, i_c, 1, 1] = 1
        self.conv.weight = torch.nn.Parameter(
            conv_weight, requires_grad=False
        )
    def forward(self, x):
        x = self.conv(x)
        out = self.depth2space(x)
        return out
if 
name
 == "
__main__
":
    bs = 2
    input_channel = 2
    h, w = 120, 150
    scale_factor = 8
    model_inputs = torch.randn(bs, input_channel, h, w)
    model = Conv2DInterpolate(
        inputs_channel=input_channel,
        scale_factor = scale_factor,
    )
    out_model = model(model_inputs)
    out_func = F.interpolate(
        model_inputs,
        scale_factor=scale_factor,
        mode="nearest"
    )
    print((out_model - out_func).max())
    print((out_model - out_func).min())

这段代码实现了一个自定义的卷积神经网络模块 Conv2DInterpolate，其主要目的是通过卷积和像素重排 （PixelShuffle） 操作实现图像的上采样。下面逐步解释代码的各个部分和其作用：

1. 类 Conv2DInterpolate 的定义

• 构造函数 （init）：
  • inputs_channel：输入图像的通道数，默认值为 1。
  • scale_factor：上采样的倍数，默认值为 2。
  • 该类首先通过 nn.Conv2d 创建一个卷积层 conv，其输入通道数为 inputs_channel，输出通道数是 inputs_channel * (scale_factor^2)，这个设计是为了后续进行像素重排时所需的通道数量。
  • 卷积的核大小是 3x3，且没有偏置 （bias=False），使用填充 1 （padding=1），以保持输入和输出的空间尺寸一致。
  • depth2space：使用了 PyTorch 中的 PixelShuffle 层进行像素重排（像素块的重新排列）。这里的 scale_factor 控制着上采样的倍数，目的是将卷积结果的通道数重排为一个更高分辨率的图像。
• 权重初始化函数 _init_weights：
  • 初始化卷积层的权重。此函数将卷积层的权重初始化为零，并根据一定规则修改权重。具体来说，它设置卷积核的中心位置 （conv_weight[i_N, i_c, 1, 1] = 1），以确保通过卷积操作得到期望的像素值。
  • 这个操作的目的是使得卷积层在初始化时生成一个有意义的初始权重，从而为后续的像素重排操作提供有效的输入。

2. 前向传播函数 forward：

• 对输入张量 x 先进行一次卷积操作 self.conv(x)，然后通过 self.depth2space(x) 进行像素重排 （PixelShuffle），最终实现图像的上采样。此操作将通道数较高的特征图转换为空间分辨率较大的输出。

3. 主程序部分：

• 在 if name == "__main__": 代码块中：
  • 创建了一个随机的输入张量 model_inputs，大小为 (batch_size=2, input_channel=2, height=120, width=150)。
  • 然后实例化了 Conv2DInterpolate 模型并进行前向传播。
  • 另外，还使用了 F.interpolate 进行基于最近邻插值的上采样操作，作为对比。

4. 输出差异对比：

通过 out_model - out_func，计算自定义模型的输出 （out_model） 与 F.interpolate 结果 （out_func） 之间的差异，并打印出最大值和最小值。

• 这部分的目的是验证自定义的 Conv2DInterpolate 模型是否能与 F.interpolate 的最近邻插值方法产生相似的结果。如果两者结果的差异很小，说明自定义模型的实现效果与标准的上采样方法接近。

tensor(0.) 
tensor(0.)

3.注意事项

若放大系数大于 2，建议使用多组 conv+depth2space 代替 resize，以实现较好的性能。根据实测经验，若 resize 放大系数为 8，且只使用一组 conv+depth2space 做 8 倍上采样时，板端运行效率很差，而使用 3 组 2 倍上采样的 conv+depth2space，板端运行耗时会回到合理范围。

此外，对于 征程 6 平台，ConvTranspose 支持了 int16 量化精度（征程 5 不支持），因此也可以考虑使用 ConvTranspose 替代 Resize（需要重训模型）。