pytorch转onnx, onnx 12 中没有hardswish opt

解决ONNX转换模型中Hardswish激活函数不支持问题

原创已于 2022-04-09 17:29:39 修改 · 2.4k 阅读

6 ·

CC 4.0 BY-SA版权

文章标签：

#onnx #hardswish #pytorch

于 2022-04-09 16:28:04 首次发布

Python基础同时被 3 个专栏收录

16 篇文章

订阅专栏

Pytorch源码学习

11 篇文章

订阅专栏

小技巧

7 篇文章

订阅专栏

该博客介绍了在ONNXOpset12下转换GhostNet、MobileNetv3Small、EfficientNetLite0和PP-LCNet等模型时遇到的Hardswish激活函数不被支持的问题。解决方案是自定义Hardswish类，并遍历模型替换所有nn.Hardswish层。示例代码给出了如何在PP-LCNet模型中进行替换，并提供了两种检查和替换Hardswish层的方法。

在onnx opset 12下转以下模型时因不支持hardswish激活函数而报错

GhostNet
MobileNetv3Small
EfficientNetLite0
PP-LCNet
解决方案是找到对应的nn.Hardswish层，将其替换为自己覆写的Hardswish实现:

class Hardswish(nn.Module):  # export-friendly version of nn.Hardswish()
    @staticmethod
    def forward(x):
        # return x * F.hardsigmoid(x)  # for torchscript and CoreML
        return x * F.hardtanh(x + 3, 0., 6.) / 6.  # for torchscript, CoreML and ONNX

以PP-LCNet为例，找到哪些层是Hardswish层，替换方法为

# 替换函数, 参考https://zhuanlan.zhihu.com/p/356273702
def _set_module(model, submodule_key, module):
    tokens = submodule_key.split('.')
    sub_tokens = tokens[:-1]
    cur_mod = model
    for s in sub_tokens:
        cur_mod = getattr(cur_mod, s)
    setattr(cur_mod, tokens[-1], module)
for k, m in model.named_modules():
	if 'dw_sp.2' in k or 'dw_sp.6' in k:
		_set_module(model, k, Hardswish())

当然也可以根据m来判断是否为nn.Hardswish的实例，

for k, m in model.named_modules():
	if isinstance(m, nn.Hardswish):
		_set_module(model, k, Hardswish())

参考

YOLOv5-Multibackbone-Compression
Pytorch替换model对象任意层的方法