TL;DR:onnxruntime
如何使用model.whatever_function(input)
而不是model.forward(input)
?我使用CLIP嵌入为我的图像和文本创建嵌入:
代码来自官方git merge
! pip install ftfy regex tqdm
! pip install git+https://github.com/openai/CLIP.git
import clip
import torch
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model, preprocess = clip.load("RN50", device=device) # Load any model
model = model.eval() # Inference Only
img_size = model.visual.input_resolution
dummy_image = torch.randn(10, 3, img_size, img_size).to(device)
image_embedding = model.encode_image(dummy_image).to(device))
dummy_texts = clip.tokenize(["quick brown fox", "lorem ipsum"]).to(device)
model.encode_text(dummy_texts)
,它工作得很好,给我[Batch, 1024]
张量对于加载的模型。
现在我已经在Onnx
中将我的模型量化为:
model.forward(dummy_image,dummy_texts) # Original CLIP result (1)
torch.onnx.export(model, (dummy_image, dummy_texts), "model.onnx", export_params=True,
input_names=["IMAGE", "TEXT"],
output_names=["LOGITS_PER_IMAGE", "LOGITS_PER_TEXT"],
opset_version=14,
dynamic_axes={
"IMAGE": {
0: "image_batch_size",
},
"TEXT": {
0: "text_batch_size",
},
"LOGITS_PER_IMAGE": {
0: "image_batch_size",
1: "text_batch_size",
},
"LOGITS_PER_TEXT": {
0: "text_batch_size",
1: "image_batch_size",
},
}
)
,模型保存。
当我测试模型为:
# Now run onnxruntime to verify
import onnxruntime as ort
ort_sess = ort.InferenceSession("model.onnx")
result=ort_sess.run(["LOGITS_PER_IMAGE", "LOGITS_PER_TEXT"],
{"IMAGE": dummy_image.numpy(), "TEXT": dummy_texts.numpy()})
它给了我一个长度为2的列表,每个图像和文本一个,result[0]
的形状为[Batch,2]
。
如果模块上的encode_image
没有调用forward
,那么在导出到Onnx之前没有什么可以阻止您覆盖forward
:
>>> model.forward = model.encode_image
>>> torch.onnx.export(model, (dummy_image, dummy_texts), "model.onnx", ...))