不同维度转onnx的问题
#1
by
cseeeee
- opened
请问可以将不同的维度转为onnx吗?尝试使用optimum将4096转为onnx,但生成的向量维度还是768.
不要用optimum,不好用,自己手动转一下,参考代码:
import torch
import torch.nn as nn
from os.path import join
from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModel
class VECModelWrapper(nn.Module):
def __init__(self, model_dir: str):
super().__init__()
self.model = AutoModel.from_pretrained(model_dir).cpu()
t = AutoTokenizer.from_pretrained(model_dir)
print(t([["123", "4,5,6"], ["123", "4,5,6"]]).keys())
# for k,v in self.model.named_parameters():
# print(k,v.shape)
self.vector_linear = torch.nn.Linear(in_features=768, out_features=4096)
vector_linear_dict = {
k.replace("linear.", ""): v for k, v in
torch.load(join(model_dir, f"2_Dense_4096/pytorch_model.bin"), map_location="cpu").items()
}
self.vector_linear.load_state_dict(vector_linear_dict)
self.model.eval()
self.vector_linear.eval()
def forward(
self,
input_ids,
# token_type_ids,
attention_mask
):
last_hidden_state = self.model(
input_ids=input_ids,
# token_type_ids=token_type_ids,
attention_mask=attention_mask
)[0]
last_hidden = last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0)
vectors = last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
vectors = self.vector_linear(vectors)
print(vectors.shape)
print(vectors[:, :4])
return vectors
def convert_pt_to_onnx(model_dir="./output/tiny_bert/best-f1", save_path="sim_model.onnx"):
# 加载模型
model = VECModelWrapper(model_dir=model_dir)
# 转为onnx
# 要准备一个输入以便跟踪模型结构,随意一个输入就行
input_ids = torch.tensor([[1, 2, 3], [11, 12, 113]], dtype=torch.long)
attention_mask = torch.tensor([[1, 1, 1], [1, 1, 0]], dtype=torch.long)
# token_type_ids = torch.tensor([[0, 0, 0], [0, 1, 0]], dtype=torch.long)
torch.onnx.export(model, # 模型
args=(
input_ids,
attention_mask,
# token_type_ids
), # 模型输入
f=save_path, # 存储路径
export_params=True,
opset_version=14, # op版本,越高越好
do_constant_folding=True,
input_names=[
'input_ids',
"attention_mask",
# "token_type_ids"
], # 模型输入的名字,随便写
output_names=['vectors'], # 模型输出的名字,使用ort就要根据这个名字获取结果
dynamic_axes={"input_ids": [0, 1], # 动态维度,第0和1维度都是可变的
"attention_mask": [0, 1],
# "token_type_ids": [0, 1],
"vectors": [0],
})
if __name__ == "__main__":
model_dir = r"E:\pt\puff-base-v1"
# model_dir = r"E:\PublicModels\bge-small-zh-v1.5"
onnx_path = r"E:\onnx\puff-base-v1\1\model.onnx"
convert_pt_to_onnx(model_dir=model_dir, save_path=onnx_path)
谢谢!
cseeeee
changed discussion status to
closed