导出ONNX模型

简介

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

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

用PyTorch框架导出ONNX模型时，框架中ACL_OP_SELECT_IMPL_MODE选项默认值为“high_precision”，ACL_OP_SELECT_IMPL_MODE详情可参见《CANN 应用开发接口》的“aclCompileOpt”章节。用户在使用导出的ONNX模型进行模型转换时，可参见《CANN ATC离线模型编译工具用户指南》中的“--op_select_implmode”章节设置与训练时相同的模式，以避免因模式选择不同而出现的精度或者性能差异。

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

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

import torch
import torch.nn as nn
import torch.onnx
import torch_npu
from torch_npu.contrib import transfer_to_npu

device = torch.device("cuda")


class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.weight = nn.Parameter(torch.randn(20, 10))
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))

def convert():
    model = ToyModel()  
    model.load_state_dict(torch.load('state_dict_model.pt', map_location='cuda'),strict=False)    #根据实际文件名称修改

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

    dummy_input = torch.randn(20, 10) #  定义输入shape
    torch.onnx.export(model, 
                      dummy_input, 
                      "model.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模型时需要先确定保存时的信息，有时保存的节点名称和模型定义中的节点会有差异，例如会多出前缀和后缀。在进行转换的时候，可以对节点名称进行修改。转换代码样例如下：

import torch
import torch.nn as nn
import torch.onnx
import torch_npu
from torch_npu.contrib import transfer_to_npu
from collections import OrderedDict

device = torch.device("cuda")

class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.weight = nn.Parameter(torch.randn(20, 10))
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))

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():
    model = ToyModel()
    checkpoint = torch.load('checkpoint.pth.tar', map_location='cuda')  # 根据实际文件名称修改
    checkpoint['state_dict'] = proc_nodes_module(checkpoint, 'state_dict')
    model.load_state_dict(checkpoint, strict=False)

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

    dummy_input = torch.randn(20, 10) #  定义输入shape
    torch.onnx.export(model,
                      dummy_input,
                      "model.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模型

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