极智开发 | 自定义结构/算子导出onnx模型

大家好，我是极智视界，本文介绍一下 自定义结构/算子导出onnx模型。

在深度学习中，把 pytorch 模型直接导出为 onnx 模型是一个很自然的行为，但是有时候也有一些比较细的需求，比如要把模型中的一个片段、一个结构甚至是一个算子导出为单独的 onnx，以方便调试。这样的话，应该怎么弄呢。

比如这里我举个例子，想要将下面这个函数导出保存为 onnx 模型：

def normalize(x, axis=-1):
    x = 1. * x / (torch.norm(x, 2, axis, keepdim=True).expand_as(x) + 1e-12)
    return x

首先它本是一个函数，如果放在深度学习模型中，可以说是几个算子组成的结构，里面包括了五个基础算子，分别是 乘、除、torch.norm、expand_as 还有加：

• torch.mul；
• torch.div；
• torch.norm;
• expand_as;
• torch.add

函数本身是不能直接导出为 onnx 模型，所以我们要把它改写为模型结构定义的方式，也就是 class 类的形式，如下改造：

class Normalize(nn.Module):
    def __init__(self):
        super(Normalize, self).__init__()

    def forward(self, x):
        x = 1. * x / torch.norm(x, 2, -1, keepdim=True).expand_as(x) + 1e-12
        return x

然后就可以通过经典的 onnx 导出方式导出来：

# # export onnx model
# 先实例化
model = Normalize()
nor_data = torch.ones((32, 512))
torch.onnx.export(model, nor_data, "./norm.onnx", verbose=True, opset_version=11)
feat_list = normalize(feat_list)

通过上面的代码就能导出为 norm.onnx 模型了，来看一下它的结构：

可以看到结构还是挺清晰的，也是可以和代码对应起来的。

而且，这种方式也适应其他小片段甚至是单算子的 onnx 导出，这里不多作赘述。

好了，以上分享了 自定义结构/算子导出onnx模型，希望我的分享能对你的学习有一点帮助。

【极智视界】

《极智开发 | 自定义结构/算子导出onnx模型》

畅享人工智能的科技魅力，让好玩的AI项目不难玩。邀请您加入我的知识星球，星球内我精心整备了大量好玩的AI项目，皆以工程源码形式开放使用，涵盖人脸、检测、分割、多模态、AIGC、自动驾驶、工业等。不敢说会对你学习有所帮助，但一定非常好玩，并持续更新更加有趣的项目。