torch模型转onnx

加载模型

model=torch.load('saved_model/moudle_best_auc.pth', map_location='cpu')
model.eval().cpu()

注：由于导出的模型是用于推理的，因此必须指定模型加载的位置和模型验证的位置，这里我使用了cpu做出导出的硬件

分析模型的输入和输出

这里要看自己的代码，在数据加载的时候输出是什么样子的，比如我这里：

那输入的list就比较好确定了，是["input_ids","input_mask"]，模型的输出list确定，要根据网络的最后一层来看，我这里用的class_loj，那我的输出list就是["class_loj"]。

构建一个dummy_input

这是导出是否成功的核心之一，这里给出了模型的输入维度，要求输入是tensor，在我的模型里面，输入的是整数的tensor，input_ids和input_mask都经过了tokenizer的编译，所以这里用的dummy_input一定也要这样子：

input_ids=np.random.randint(10, size=(batch_size, 373))# 373是维度，根据自己的来定，这里是序列问题
inputs = torch.from_numpy(tape_pad(input_ids, 0))#我需要将其转为这样的格式
inputs_mask = torch.from_numpy(tape_pad(np.ones_like(input_ids), 0)) # 构建mask的方式根据自己选择
dummy_input={"input_ids":inputs,
            "input_mask":input_mask} # 得到这样一个dummy_input

导出模型

这些全部准备好了之后，就可以导出了：

torch.onnx.export(
    model, # 前面torch.load导入的模型
    (dummy_input,),#模型的输入
    'tape.onnx',  # 导出的模型名字
    export_params=True, # 如果指定为True或默认, 参数也会被导出. 如果你要导出一个没训练过的就设为 False.
    opset_version=15, # 版本号，写10-15都可以
    do_constant_folding=True, # 是否执行常量折叠优化
    input_names=model_inputs, # 输入模型的张量的名称
    output_names=model_outputs, # 模型输出的张量的名称
    # dynamic_axes将batch_size的维度指定为动态，
    dynamic_axes={'input_ids': {0: 'batch_size', 1: 'sequence'},
                  'input_mask':{0: 'batch_size', 1: 'sequence'},
                  'class_loj': {0: 'batch_size'}}
    
)

完整导出代码

import torch
from tape.tokenizers import TAPETokenizer
from tape.datasets import pad_sequences as tape_pad
import numpy as np

# 导入token和模型
tokenizer=TAPETokenizer(vocab='unirep')
model=torch.load('saved_model/moudle_best.pth', map_location='cpu')
model.eval().cpu()

# 准备输入格式
batch_size=1
input_ids=np.random.randint(10, size=(batch_size, 373))
inputs = torch.from_numpy(tape_pad(input_ids, 0))
inputs_mask = torch.from_numpy(tape_pad(np.ones_like(input_ids), 0))
dummy_input={"input_ids":inputs,
            "input_mask":input_mask}

# 准备输入和输出张量名称            
model_inputs=['input_ids', 'input_mask']
model_outputs=['class_loj']

# 使用onnx导出
torch.onnx.export(
    model,
    (dummy_input,),
    'tape.onnx',
    export_params=True,
    opset_version=15,
    do_constant_folding=True,
    input_names=model_inputs,
    output_names=model_outputs,
    dynamic_axes={'input_ids': {0: 'batch_size', 1: 'sequence'},
                  'input_mask':{0: 'batch_size', 1: 'sequence'},
                  'class_loj': {0: 'batch_size'}}
    
)