解决Failed to export an ONNX attribute ‘onnx::Gather‘ 报错

最新推荐文章于 2025-06-09 16:16:43 发布

梦坠凡尘(AICV与前沿)

最新推荐文章于 2025-06-09 16:16:43 发布

阅读量2.1k

点赞数 9

CC 4.0 BY-SA版权

分类专栏：深度学习 Pytorch 文章标签：计算机视觉深度学习 Pytorch

本文链接：https://blog.youkuaiyun.com/c2250645962/article/details/114700033

深度学习同时被 2 个专栏收录

74 篇文章

订阅专栏

Pytorch

16 篇文章

订阅专栏

本文介绍了解决使用CAAttention模块进行ONNX转换时遇到的AdaptiveAvgPool2d参数非常量化问题的方法。通过将池化层的尺寸参数在运行时动态获取并设置为常量值的方式，成功实现了模型的ONNX转换。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

问题

使用CAAttention转换onnx时出现如题错误，原因是AdaptiveAvgPool2d()里面的参数没有常量化。

class CoordAtt(nn.Module):
    def __init__(self, inp, oup, reduction=32):
        super(CoordAtt, self).__init__()
        self.pool_h = nn.AdaptiveAvgPool2d((None, 1))
        self.pool_w = nn.AdaptiveAvgPool2d((1, None))

        mip = max(8, inp // reduction)

        self.conv1 = nn.Conv2d(inp, mip, kernel_size=1, stride=1, padding=0)
        self.bn1 = nn.BatchNorm2d(mip)
        self.act = h_swish()

        self.conv_h = nn.Conv2d(mip, oup, kernel_size=1, stride=1, padding=0)
        self.conv_w = nn.Conv2d(mip, oup, kernel_size=1, stride=1, padding=0)


    def forward(self, x):
        identity = x

        n,c,h,w = x.size()
       
        x_h = self.pool_h(x)
        x_w = self.pool_w(x).permute(0, 1, 3, 2)

        y = torch.cat([x_h, x_w], dim=2)
        y = self.conv1(y)
        y = self.bn1(y)
        y = self.act(y)
        x_h, x_w = torch.split(y, [h, w], dim=2)
        x_w = x_w.permute(0, 1, 3, 2)
        a_h = self.conv_h(x_h).sigmoid()
        a_w = self.conv_w(x_w).sigmoid()
        out = identity * a_w * a_h

        return out

解决方案

修改代码如下：

class CoordAtt(nn.Module):
    def __init__(self, inp, oup, reduction=32):
        super(CoordAtt, self).__init__()
        # self.pool_h = nn.AdaptiveAvgPool2d((None, 1))
        # self.pool_w = nn.AdaptiveAvgPool2d((1, None))
        mip = max(8, inp // reduction)

        self.conv1 = nn.Conv2d(inp, mip, kernel_size=1, stride=1, padding=0)
        self.bn1 = nn.BatchNorm2d(mip)
        self.act = h_swish()

        self.conv_h = nn.Conv2d(mip, oup, kernel_size=1, stride=1, padding=0)
        self.conv_w = nn.Conv2d(mip, oup, kernel_size=1, stride=1, padding=0)


    def forward(self, x):
        identity = x
        n,c,h,w = x.size()

        if torch.is_tensor(h):
            h = h.item()  # 这里是修正代码
            w = w.item()  # 这里是修正代码
        pool_h = nn.AdaptiveAvgPool2d((h, 1))
        pool_w = nn.AdaptiveAvgPool2d((1, w))

        x_h = pool_h(x)
        x_w = pool_w(x).permute(0, 1, 3, 2)

        y = torch.cat([x_h, x_w], dim=2)
        y = self.conv1(y)
        y = self.bn1(y)
        y = self.act(y)
        x_h, x_w = torch.split(y, [h, w], dim=2)
        x_w = x_w.permute(0, 1, 3, 2)
        a_h = self.conv_h(x_h).sigmoid()
        a_w = self.conv_w(x_w).sigmoid()

        out = identity * a_w * a_h

        return out