1.技術(shù)背景

在進(jìn)行模型壓縮與加速的過(guò)程中，量化技術(shù)成為了提升推理速度和降低計(jì)算資源消耗的重要手段。然而，在實(shí)際應(yīng)用中，許多用戶發(fā)現(xiàn)，采用 resize 操作時(shí)，僅能使用 int8 精度量化，這一限制導(dǎo)致了模型精度的顯著下降。盡管 int8 精度在提升計(jì)算效率方面具有優(yōu)勢(shì)，但精度的喪失卻使得模型的推理結(jié)果偏差增大，給實(shí)際應(yīng)用帶來(lái)了不少困擾。如何在保證性能的同時(shí)，最大程度地減少精度損失，成為了當(dāng)前技術(shù)實(shí)現(xiàn)中的一個(gè)難題。 在當(dāng)前工具鏈版本下（J6 3.0.31），resize 算子僅支持 int8 量化精度，不支持 int16，因此該類算子有一定概率觸發(fā)精度下降問(wèn)題。

針對(duì)這一問(wèn)題，本文將介紹一種新的方法，可以有效提升上采樣操作中的精度，解決傳統(tǒng) int8 精度量化帶來(lái)的精度下降問(wèn)題。通過(guò)巧妙地優(yōu)化模型結(jié)構(gòu)，能夠在不顯著影響計(jì)算效率的前提下，顯著提高上采樣的精度。 本文介紹的這種解決方案，在不影響模型權(quán)重（無(wú)需重訓(xùn)）的情況下，通過(guò)算子替換，使得上采樣功能支持 int16 量化精度，以解決精度下降問(wèn)題。 相信這一方案將為廣大開(kāi)發(fā)者帶來(lái)幫助～

2.方案介紹

onnx 中的 resize 算子，在 pytorch 代碼中常表現(xiàn)為 F.interpolate 函數(shù)。當(dāng) F.interpolate 的 mode 為 nearest 時(shí)，該函數(shù)的功能和 conv+depth2space 完全等效。而 conv 和 depth2space 均支持 int16 量化，因此可以通過(guò)算子替換的方式變相實(shí)現(xiàn)上采樣的 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())

這段代碼實(shí)現(xiàn)了一個(gè)自定義的卷積神經(jīng)網(wǎng)絡(luò)模塊 Conv2DInterpolate，其主要目的是通過(guò)卷積和像素重排 （PixelShuffle） 操作實(shí)現(xiàn)圖像的上采樣。下面逐步解釋代碼的各個(gè)部分和其作用：

1. 類 Conv2DInterpolate 的定義

構(gòu)造函數(shù) （init）：

inputs_channel：輸入圖像的通道數(shù)，默認(rèn)值為 1。

scale_factor：上采樣的倍數(shù)，默認(rèn)值為 2。

該類首先通過(guò) nn.Conv2d 創(chuàng)建一個(gè)卷積層 conv，其輸入通道數(shù)為 inputs_channel，輸出通道數(shù)是 inputs_channel * (scale_factor^2)，這個(gè)設(shè)計(jì)是為了后續(xù)進(jìn)行像素重排時(shí)所需的通道數(shù)量。

卷積的核大小是 3x3，且沒(méi)有偏置 （bias=False），使用填充 1 （padding=1），以保持輸入和輸出的空間尺寸一致。

depth2space：使用了 PyTorch 中的 PixelShuffle 層進(jìn)行像素重排（像素塊的重新排列）。這里的 scale_factor 控制著上采樣的倍數(shù)，目的是將卷積結(jié)果的通道數(shù)重排為一個(gè)更高分辨率的圖像。

權(quán)重初始化函數(shù) _init_weights：

初始化卷積層的權(quán)重。此函數(shù)將卷積層的權(quán)重初始化為零，并根據(jù)一定規(guī)則修改權(quán)重。具體來(lái)說(shuō)，它設(shè)置卷積核的中心位置 （conv_weight[i_N, i_c, 1, 1] = 1），以確保通過(guò)卷積操作得到期望的像素值。

這個(gè)操作的目的是使得卷積層在初始化時(shí)生成一個(gè)有意義的初始權(quán)重，從而為后續(xù)的像素重排操作提供有效的輸入。

2. 前向傳播函數(shù) forward：

對(duì)輸入張量 x 先進(jìn)行一次卷積操作 self.conv(x)，然后通過(guò) self.depth2space(x) 進(jìn)行像素重排 （PixelShuffle），最終實(shí)現(xiàn)圖像的上采樣。此操作將通道數(shù)較高的特征圖轉(zhuǎn)換為空間分辨率較大的輸出。

3. 主程序部分：

在 if name == "__main__": 代碼塊中：

創(chuàng)建了一個(gè)隨機(jī)的輸入張量 model_inputs，大小為 (batch_size=2, input_channel=2, height=120, width=150)。

然后實(shí)例化了 Conv2DInterpolate 模型并進(jìn)行前向傳播。

另外，還使用了 F.interpolate 進(jìn)行基于最近鄰插值的上采樣操作，作為對(duì)比。

4. 輸出差異對(duì)比：

 通過(guò) out_model - out_func，計(jì)算自定義模型的輸出 （out_model） 與 F.interpolate 結(jié)果 （out_func） 之間的差異，并打印出最大值和最小值。

這部分的目的是驗(yàn)證自定義的 Conv2DInterpolate 模型是否能與 F.interpolate 的最近鄰插值方法產(chǎn)生相似的結(jié)果。如果兩者結(jié)果的差異很小，說(shuō)明自定義模型的實(shí)現(xiàn)效果與標(biāo)準(zhǔn)的上采樣方法接近。

tensor(0.) 
tensor(0.)

3.注意事項(xiàng)

若放大系數(shù)大于 2，建議使用多組 conv+depth2space 代替 resize，以實(shí)現(xiàn)較好的性能。根據(jù)實(shí)測(cè)經(jīng)驗(yàn)，若 resize 放大系數(shù)為 8，且只使用一組 conv+depth2space 做 8 倍上采樣時(shí)，板端運(yùn)行效率很差，而使用 3 組 2 倍上采樣的 conv+depth2space，板端運(yùn)行耗時(shí)會(huì)回到合理范圍。

此外，對(duì)于 征程 6 平臺(tái)，ConvTranspose 支持了 int16 量化精度（征程 5 不支持），因此也可以考慮使用 ConvTranspose 替代 Resize（需要重訓(xùn)模型）。