PyTorch转ONNX踩坑——RuntimeError: ONNX export of operator adaptive_avg_pool2d & Failed to export an ONNX attribute 'onnx::Gather'

2024年11月8日 317点热度 1人点赞 0条评论

踩坑一：

RuntimeError: ONNX export of operator adaptive_avg_pool2d

报错原因：这个错误表明在将模型导出为 ONNX 格式时，adaptive_avg_pool2d 操作不支持当前的输出大小，因为输出大小不是输入大小的因数。因为onnx要求模型的输入参数的固定的，而自适应池化层是根据输入来确定模型参数的。
ChatGPT提供了一些别的方法，如更新Pytorch版本或尝试使用 ONNX opset 版本或 dynamic_axes 参数，经过测试都无法解决以上问题。
只能采用：

1. 自定义 adaptive_avg_pool2d函数代替
2. 修改模型代码以使用 avg_pool2d代替

原理：参考

stride = floor ( (input_size / (output_size) )
kernel_size = input_size − (output_size−1) * stride
padding = 0

具体修改：
1、用nn.AvgPool2d替换nn.AdaptiveAvgPool2d(1)

###先定义好要替换的类：
class CustomGlobalAvgPool2D(nn.Module):
    def __init__(self, input_size):
        super(CustomGlobalAvgPool2D, self).__init__()
        self.pool = nn.AvgPool2d(kernel_size=input_size)

    def forward(self, x):
        return self.pool(x)

###需要指定输入特征的size，具体使用如下：
def forward(self, x):
        x128,x64,x32,x16 = self.backbone(x)

        #### directional Prior ####
        directional_c5 = self.channel_mapping(x16)
        mapped_c5=F.interpolate(directional_c5,scale_factor=32,mode='bilinear',align_corners=True)
        mapped_c5 = self.direc_reencode(mapped_c5)   # torch.Size([10, 24, 512, 512])
        self.map_avg_pool = CustomGlobalAvgPool2D(input_size=[512, 512])
        d_prior = self.map_avg_pool(mapped_c5)   # torch.Size([10, 24, 1, 1])

2、用nn.AvgPool2d替换nn.adaptive_avg_pool2d

###原代码内容：
def forward(self, x):
        batch_size, C, H, W = x.shape
        x_pooled = x

        if H != self.dct_h or W != self.dct_w:
            x_pooled = torch.nn.functional.adaptive_avg_pool2d(x, (self.dct_h, self.dct_w))   ###被指定了输出特征的size

###定义要替换的类：
class CustomAvgPool2D(nn.Module):
    def __init__(self, input_size, output_size):
        super(CustomAvgPool2D, self).__init__()
        self.input_size = input_size
        self.output_size = output_size

        # 使用上采样层
        self.upsample = nn.Upsample(size=output_size, mode='bilinear', align_corners=False)

    def forward(self, x):
        _, _, h, w = x.shape
        dct_h, dct_w = self.output_size
        if torch.is_tensor(h):

        # 如果输入尺寸小于目标尺寸，则使用上采样。
        # 这是由于我自己的模型中存在此情况，需要额外处理。没有该情况的可以直接执行else的内容。
        if h < dct_h or w < dct_w:
            return self.upsample(x)
        else:
            # 否则直接使用 AvgPool2d
            stride = (max(1, h // dct_h), max(1, w // dct_w))
            kernel_size = (h - (dct_h - 1) * stride[0], w - (dct_w - 1) * stride[1])
            pool = nn.AvgPool2d(kernel_size=kernel_size, stride=stride, padding=0)
            return pool(x)

###修改后的代码：
def forward(self, x):
        batch_size, C, H, W = x.shape
        x_pooled = x

        if H != self.dct_h or W != self.dct_w:
            self.avg_pool = CustomAvgPool2D(input_size=(H, W), output_size=(self.dct_h, self.dct_w))   #需要确定输入size和输出size，保证一致
            x_pooled = self.avg_pool(x)

以上操作即可保证处理结果一致。

踩坑二：

Failed to export an ONNX attribute 'onnx::Gather'
修改完网络中的全部nn.AdaptiveAvgPool2d(1)函数和nn.adaptive_avg_pool2d（）函数，执行转ONNX还会出现以上错误！！！泪目！！！

修改方法参考
H与 W的类型会变成 torch.tensor, 使得转换报错，所以需要在转换前加上两行代码:

###修改一：
class CustomAvgPool2D(nn.Module):
    def __init__(self, input_size, output_size):
        super(CustomAvgPool2D, self).__init__()
        self.input_size = input_size
        self.output_size = output_size

        # 使用上采样层
        self.upsample = nn.Upsample(size=output_size, mode='bilinear', align_corners=False)

    def forward(self, x):
        _, _, h, w = x.shape
        dct_h, dct_w = self.output_size
        if torch.is_tensor(h):
            h = h.item()  # 这里是修正代码
            w = w.item()  # 这里是修正代码

        # 如果输入尺寸小于目标尺寸，则使用上采样
        if h < dct_h or w < dct_w:
            return self.upsample(x)
        else:
            # 否则直接使用 AvgPool2d
            stride = (max(1, h // dct_h), max(1, w // dct_w))
            kernel_size = (h - (dct_h - 1) * stride[0], w - (dct_w - 1) * stride[1])
            pool = nn.AvgPool2d(kernel_size=kernel_size, stride=stride, padding=0)
            return pool(x)

###修改二：
def forward(self, x):
    x_pooled_spectral = x_pooled * params
     xB, xC, xH, xW = x_pooled_spectral.shape
      if torch.is_tensor(xH):
            xH = xH.item()  # 这里是修正代码
            xW = xW.item()  # 这里是修正代码
    self.avg_pool_1 = CustomGlobalAvgPool2D(input_size=(xH, xW))
     multi_spectral_feature_avg += self.avg_pool_1(x_pooled_spectral)

Post Views: 321