PyTorch模型转onnx格式

因项目需要，把PyTorch模型转为onnx。

Huggingface的模型无法直接转为onnx（仅有部分支持），如果是huggingface的模型，需要先转成PyTorch，详见文章《Huggingface模型转PyTorch模型》。

onnx是一种通用的模型架构，详细介绍可参考官网。

转换代码：

from transformers import AutoTokenizer
from transformers import AutoModel
from transformers import LlamaTokenizer, LlamaForCausalLM
import torch

# convert pytorch model to onnx model

def generate_prompt(text):
    return f"""Below is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
{text}

### Response:"""


text = '请介绍一下中国海洋大学'
prompt = generate_prompt(text)
tokenizer = LlamaTokenizer.from_pretrained('minlik/chinese-alpaca-7b-merged')
input = tokenizer.encode(prompt, return_tensors='pt')
print("now load pytorch model")
modeloftorchtype = torch.load('your_model_path_and_name.pt')
print("load finished, now eval")
modeloftorchtype.eval()
print("eval finished, now convert to onnx")

torch.onnx.export(modeloftorchtype, input, "your_model_path_and_name.onnx",input_names=["input"], output_names=["output"],opset_version=11,dynamic_axes={'input': {1:'tokenlength'},'output' : [0,1]}) # 输入的第1为可变

转换的思路其实很简单，就是把原有的模型，推理一遍，就可以使用PyTorch自带的函数导出了。

这里面有几点需要注意：

1. 转换前需要调用eval函数
2. 如果模型很大，那么转换出来的文件会很多，因为存储onnx的文件大小不允许超过2G，这是正常的
3. export函数中，input_names=["input"], output_names=["output"]是自己定义的。dynamic_axes={'input': {1:'tokenlength'},'output' : [0,1]}是用来定义输入输出的可变长度的。比如，对于文本问答模型，每次输入的长度都不统一，如果按照这次转换的推理，那么input的维度就固定了，这样当再次加载模型推理的时候，就会报长度不一致的错误，所以这里要设置为可变长度。其中，1表示三维张量的第二维是可变的。'output' : [0,1]表示output张量（列表）的第一维和第二维，都是可变的。