pytorch中bin模型文件转onnx遇到的问题
1 常规做法
import os
import numpy as np
from transformers import GPT2LMHeadModel
import torch
localfile = r"C:\Users\min_ppl_model"
model = GPT2LMHeadModel.from_pretrained(localfile)
# 输入shape 为 1，50 其中1 为bs 50为固定输入长度，前面做量化时规定的。
# 确定输入shape
batch_size = 1
seq = 50
example_input = torch.randn(batch_size, seq)
print(example_input)
# 保存文件
save_onnx_dir = r"layer20_gpt2"
os.makedirs(save_onnx_dir, exist_ok=True)
save_onnx_model = r"layer20_gpt2\layer20_gpt2.onnx"
# 转onnx
with torch.no_grad():
torch.onnx.export(model,
example_input,
save_onnx_model,
opset_version=11,
input_names=["inp"], output_names=["opt"])
print("-----end-------")
其中
example_input = torch.randn(batch_size, seq, dtype=torch.float32)
最为关键一步
类型为：
(Pdb) type(example_input)
<class 'torch.Tensor'>
打印为：
tensor([[-2.0319, -0.4021, 0.0092, 1.4100, -0.2214, 0.6954, 0.1764, 0.2111,
-0.4725, -0.7527, -0.0766, -1.4510, -0.2528, -1.4077, -0.9340, 0.2309,
0.5564, -0.0498, -0.7499, -1.8176, -0.8981, 1.3997, 0.2904, -0.5024,
0.8392, -0.2341, -0.5459, -1.0992, 0.0211, 0.3346, -0.5087, -0.6159,
1.3256, -0.0423, -0.0764, 1.0469, -1.7328, -1.3470, -0.3346, 0.2129,
1.1073, -0.7503, 0.3968, -0.1374, 0.6514, -0.8763, -1.1972, -1.7750,
-0.9977, -2.2836]])
报错为：RuntimeError: Expected tensor for argument #1 'indices' to have one of the following scalar types: Long, Int; but got torch.FloatTensor instead (while checking arguments for embedding)
原因为 embedding 为int。
直接更换为整型：
example_input = torch.randint(1,5, size=(batch_size, seq))
2 先用np 再转torch.tensor
example_input = torch.from_numpy(np.random.randint(low=6,high=10,size=(1,50),dtype=np.int32))
可以稍微指定一个整数范围
(Pdb) type(example_input)
<class 'torch.Tensor'>
打印信息：
tensor([[6, 9, 9, 9, 8, 6, 6, 6, 8, 8, 9, 9, 9, 7, 9, 9, 7, 8, 8, 7, 9, 9, 9, 6,
7, 6, 8, 8, 9, 9, 8, 7, 7, 8, 8, 9, 6, 9, 7, 9, 8, 6, 9, 7, 7, 7, 9, 6,
7, 8]], dtype=torch.int32)
则可以成功执行。
总结：
1 能先用int先用int， 要求是float就先用float， 随机生成数据。
2 如果定义数据是dtype=torch.int32 如果带了，则打印数据时也会带， 否则不显示。
tensor([[6, 9, 9, 9, 8, 6, 6, 6, 8, 8, 9, 9, 9, 7, 9, 9, 7, 8, 8, 7, 9, 9, 9, 6,
7, 6, 8, 8, 9, 9, 8, 7, 7, 8, 8, 9, 6, 9, 7, 9, 8, 6, 9, 7, 7, 7, 9, 6,
7, 8]], dtype=torch.int32)