ReXNet模型pth转onnx,swish不支持问题

问题描述:

 最近在使用一个算法模型ReXNet,算法是基于pytorch实现的,训练模型保存格式为.pth,我这里想把它转化为onnx做推理加速,模型转化代码这里就不贴了,很简单,直接上转化过程中遇到问题如下图:

可以看到报错信息为:RuntimeError: ONNX export failed: Couldn’t export Python operator SwishJitImplementation
说明SwishJitImplementation这个函数中有onnx不支持的op,我们去到rReXNet网络中去找到SwishJitImplementation这个函数,代码片如下图

即为模型中使用了swish激活函数,而onnx是不支持Swish这个函数的,我们需要对算法模型进行修改,或者在模型转化时对swish进行修改

首先先看下swish这个函数公式为:

函数图像为:

当 β \beta β = 0时，Swish激活函数变为线性函数 f(x) = x 2 \frac{x}{2} 2x​，
当 β \beta β = ∞ \infty ∞ 时，Swish激活函数变为0或x，相当于Relu，
所以，Swish函数可以看作是介于线性函数与ReLU函数之间的平滑函数

通过函数公式可以看出swish函数即为x*一个sigmoid(x)函数,sigmoid函数onnx是支持的,因此我们这里对swish进行转换,从而使pytorch模型成功转换为onnx模型

工程实现:
在Tensorflow中使用 x * tf.sigmoid(beta * x)或tf.nn.swish(x).
在Caffe中使用Scale+Sigmoid+EltWise(PROD)来实现或者合并成一个层.
在pytorch中使用x.mul_(x.sigmoid())

onnx不支持的操作还有如下:
onnx只能输出静态图，因此不支持if-else分支。一次只能走一个分支
onnx不支持步长为2的切片。例如a[::2,::2]
onnx不支持对切片对象赋值。例如a[0,:,:,:]=b， 可以用torch.cat改写
onnx里面的resize要求output shape必须为常量
不支持自定义的conv，支持conv1d,conv2d,conv3d
AdaptiveAvgPool2d : 使用AvgPool2d替换
Expand : torch.cat进行替换
ReLU6 : 通过relu改写；即 ReLU6 = 6.0 - nn.ReLU(inplace=True)(6.0-relu)
C++调用时数据类型要是32位浮点型

常用的激活函数
参考:
https://blog.csdn.net/weixin_39421413/article/details/103083806
https://blog.csdn.net/qq_38978225/article/details/124587779
https://blog.csdn.net/weixin_38649779/article/details/122037115
https://blog.csdn.net/bblingbbling/article/details/107105648