精度調優｜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（需要重訓模型）。