导出ONNX模型

简介

ONNX是业内目前比较主流的模型格式，广泛用于模型交流及部署。PyTorch模型在昇腾AI处理器上的部署策略是基于PyTorch官方支持的ONNX模块实现的。

本节主要介绍如何将Checkpoint文件通过torch.onnx.export()接口导出为ONNX模型。

使用PyTorch框架导出ONNX模型时，框架中设置算子编译选项的ACL_OP_SELECT_IMPL_MODE选项默认值为“high_precision”，用户可根据需要自行修改。用户在使用导出的ONNX模型进行模型转换时，可参见《CANN ATC工具使用指南》中的“--op_select_implmode”章节设置与训练时相同的模式，以避免因模式选择不同而出现的精度或者性能差异。

.pth或.pt文件导出ONNX模型

保存的.pth或.pt文件可以通过PyTorch构建模型，再加载权重的方法恢复，然后导出ONNX模型，样例如下：

import torch
import torch_npu
import torch.onnx
import torchvision.models as models
# 设置使用CPU导出模型
device = torch.device("cpu") 

def convert():
    # 模型定义来自于torchvision，样例生成的模型文件是基于resnet50模型
    model = models.resnet50(pretrained = False)  
    resnet50_model = torch.load('resnet50.pth', map_location='cpu')    #根据实际文件名称修改
    model.load_state_dict(resnet50_model) 

    batch_size = 1  #批处理大小
    input_shape = (3, 224, 224)   #输入数据,改成自己的输入shape

    # 模型设置为推理模式
    model.eval()

    dummy_input = torch.randn(batch_size, *input_shape) #  定义输入shape
    torch.onnx.export(model, 
                      dummy_input, 
                      "resnet50_official.onnx", 
                      input_names = ["input"],   # 构造输入名
                    output_names = ["output"],    # 构造输出名
                      opset_version=11,    # ATC工具目前支持opset_version=9，10，11，12，13
                      dynamic_axes={"input":{0:"batch_size"}, "output":{0:"batch_size"}})  #支持输出动态轴

if __name__ == "__main__":
    convert()

• 在导出ONNX模型之前，必须调用model.eval() 来将dropout和batch normalization层设置为推理模式。
• 样例脚本中的model来自于torchvision模块中的定义，用户使用自己的模型时需自行指定。
• 构造输入输出需要对应训练时的输入输出，否则无法正常推理。

.pth.tar文件导出ONNX模型

.pth.tar在导出ONNX模型时需要先确定保存时的信息，有时保存的节点名称和模型定义中的节点会有差异，例如会多出前缀和后缀。在进行转换的时候，可以对节点名称进行修改。转换代码样例如下：

from collections import OrderedDict
import torch
import torch_npu
import torch.onnx
import torchvision.models as models

# 如果发现pth.tar文件保存时节点名加了前缀或后缀，则通过遍历删除。此处以遍历删除前缀"module."为例。若无前缀后缀则不影响。
def proc_nodes_module(checkpoint, AttrName):
    new_state_dict = OrderedDict()
    for key, value in checkpoint[AttrName].items():
        if key == "module.features.0.0.weight":
            print(value)
        #根据实际前缀后缀情况修改
        if(key[0:7] == "module."):
            name = key[7:]
        else:
            name = key[0:]

        new_state_dict[name] = value
    return new_state_dict

def convert():
    # 模型定义来自于torchvision，样例生成的模型文件是基于resnet50模型
    checkpoint = torch.load("./resnet50.pth.tar", map_location=torch.device('cpu'))    #根据实际文件名称修改
    checkpoint['state_dict'] = proc_nodes_module(checkpoint,'state_dict')
    model = models.resnet50(pretrained = False)
    model.load_state_dict(checkpoint['state_dict'])
    model.eval()
    input_names = ["actual_input_1"]
    output_names = ["output1"]
    dummy_input = torch.randn(1, 3, 224, 224)
    torch.onnx.export(model, dummy_input, "resnet50.onnx", input_names = input_names, output_names = output_names, opset_version=11)    #输出文件名根据实际情况修改

if __name__ == "__main__":
    convert()

自定义算子导出ONNX模型

对于非NPU自定义算子，导出ONNX的逻辑和限制遵循PyTorch框架，请参考官方网站的PyTorch框架issue或者文档进行修改。