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为可变