极智AI | 教你简化 onnx upsample 算子

944 阅读3分钟

  一起养成写作习惯!这是我参与「掘金日新计划 · 4 月更文挑战」的第16天,点击查看活动详情

欢迎关注我的公众号 [极智视界],获取我的更多笔记分享

   大家好,我是极智视界。本文介绍了简化 onnx upsample 算子的方法。

   实际部署中经常会涉及到 pytorch / tensorflow / darknet -> onnx 的模型转换过程。本身模型转换过程就比较麻烦(当然 pytorch export onnx 十分方便),成了后你经常又会发现 convert successed 的 onnx model 看起来面目全非。这样一方面看起来不够清爽,另一方面算子细粒度太小,不利于部署。

   这里先拿转换 onnx 时的 upsample 算子简洁化说起。

1、upsample 原理

  upsample 是上采样,直观理解就是放大图像,采用各种插值算法来扩充 feature map。upsample 在分割网络、GAN 里比较常见,因为这些网络需要还原到特征图。目标检测网络常常会加入多尺度检测的特性,如在 Yolo 系列中最后的 yolo 分支往往就会考虑多尺度检测,这个时候就需要加入 upsample 算子来还原 feature map 的尺寸。

   把 upsample 看成这样就好了,当 stride = 1 时:

   当 stride = 2 时:

2、简化 onnx upsample

   这里介绍一下 pytorch -> onnx upsample 的三种实现方法,也是逐步优化的过程。

2.1 upsample expand

  看下 upsample expand 的算子实现方式:

class Upsample_expand(nn.Module):
    def __init__(self, stride=2):
        super(Upsample_expand, self).__init__()
        self.stride = stride

    def forward(self, x):
        assert (x.data.dim() == 4)
        
        x = x.view(x.size(0), x.size(1), x.size(2), 1, x.size(3), 1).\
            expand(x.size(0), x.size(1), x.size(2), self.stride, x.size(3), self.stride).contiguous().\
            view(x.size(0), x.size(1), x.size(2) * self.stride, x.size(3) * self.stride)

        return x

   来看下 upsample expand 转出来的 onnx 结构:

  以上的 onnx 结构算子细粒度太细,看起来很不清爽。

2.2 upsample interpolate

  看下 upsample interpolate 的算子实现方式:

class Upsample_interpolate(nn.Module):
    def __init__(self, stride):
        super(Upsample_interpolate, self).__init__()
        self.stride = stride

    def forward(self, x):
        assert (x.data.dim() == 4)

        x = torch.nn.functional.interpolate(x, size=(x.size(2) * self.stride, x.size(3) * self.stride), mode='nearest')
        return x

  来看下 upsample interpolate 转出来的 onnx 结构:

  相比于上面的结构是不是已经简化了很多,算子个数减少了很多,但是还不是我们最终想要的样子,让我们继续简化。

2.3 nn.Upsample

  看下 nn.Upsample 的算子实现方式:

nn.Upsample(scale_factor=stride, mode="nearest")

  来看下 nn.Upsample 转出来的 onnx 结构:

   这样是不是看起来很简洁了,nice~

   onnx 作为一个模型转换的中间结构形式,往往喜欢把算子拆成细粒度更加小的算子组合,因为这样能够通过细算子的组合形成更加多的大算子,也可以解决大算子不支持的问题。onnx 的这种特性跟部署喜欢大算子存在着天然的 Gap,所以在 onnx 和 部署之间需要我们做更多的算子融合的工作,如 yolo 层也会存在同样的问题。


  这篇先这样了,有问题欢迎交流~ 以上分享了简化 onnx upsample 算子的方法,希望我的分享能对你的学习有一点帮助。

 【公众号传送】

《【模型推理】教你简化 onnx upsample 算子》


logo_show.gif