通过onnx runtiome
核心代码是:
vision_session = ort.InferenceSession("siglip_image_embed.onnx")
image_embeds = vision_session.run(None, {input_name: pixel_values})[0]
input_name是导出模型时,指定的名字。
import numpy as np
import onnxruntime as ort
from PIL import Image
from transformers import AutoProcessor, AutoModel
import torch
# ========== 1. 加载 ONNX 会话 ==========
vision_session = ort.InferenceSession("siglip_image_embed.onnx")
text_session = ort.InferenceSession("siglip_text_embed.onnx")
# 打印输入信息,确认名称和类型
print("Vision inputs:", [(i.name, i.shape, i.type) for i in vision_session.get_inputs()])
print("Text inputs:", [(i.name, i.shape, i.type) for i in text_session.get_inputs()])
# ========== 2. 加载 processor 和原始模型(取 logit_scale / logit_bias) ==========
processor = AutoProcessor.from_pretrained("./my_local_siglip")
# 仅用于读取 logit_scale / logit_bias,不参与 ONNX 推理
ref_model = AutoModel.from_pretrained("./my_local_siglip").eval()
logit_scale = ref_model.logit_scale.exp().item()
logit_bias = ref_model.logit_bias.item()
print(f"logit_scale={logit_scale:.4f}, logit_bias={logit_bias:.4f}")
# ========== 3. 图像预处理与视觉推理 ==========
def infer_image(image_path):
image = Image.open(image_path).convert("RGB")
# processor 会按 checkpoint 的配置做 resize / normalize / to tensor
inputs = processor(images=image, return_tensors="np")
pixel_values = inputs["pixel_values"].astype(np.float32) # (1, 3, 512, 512)
input_name = vision_session.get_inputs()[0].name
image_embeds = vision_session.run(None, {input_name: pixel_values})[0]
return image_embeds
# ========== 4. 文本预处理与文本推理 ==========
def infer_text(text):
inputs = processor(
text=text,
return_tensors="np",
padding="max_length",
max_length=64,
truncation=True,
)
input_ids = inputs["input_ids"].astype(np.int64) # 与导出时 dtype 一致
input_name = text_session.get_inputs()[0].name
text_embeds = text_session.run(None, {input_name: input_ids})[0]
return text_embeds
# ========== 5. 余弦相似度与匹配概率 ==========
def cosine_similarity(a, b):
# a: (1, D), b: (1, D) -> 标量
a = a / np.linalg.norm(a, axis=-1, keepdims=True)
b = b / np.linalg.norm(b, axis=-1, keepdims=True)
return float(np.sum(a * b, axis=-1)[0])
def match_probability(cos_sim):
# SigLIP: logit = logit_scale * cos_sim + logit_bias, prob = sigmoid(logit)
logit = logit_scale * cos_sim + logit_bias
return float(1.0 / (1.0 + np.exp(-logit)))
# ========== 6. 运行示例 ==========
image_path = "2.jpg"
text = "乌龟"
image_embeds = infer_image(image_path)
text_embeds = infer_text(text)
print("image_embeds shape:", image_embeds.shape)
print("text_embeds shape:", text_embeds.shape)
cos_sim = cosine_similarity(image_embeds, text_embeds)
prob = match_probability(cos_sim)
print(f"余弦相似度: {cos_sim:.4f}")
print(f"匹配概率(sigmoid): {prob:.4f}")输出:
D:\>python sigliponnx_inf2.py
Vision inputs: [('pixel_values', [1, 3, 512, 512], 'tensor(float)')]
Text inputs: [('input_ids', [1, 64], 'tensor(int64)')]
[transformers] Model config: bos_token_id must be `None` or an integer within the vocabulary (between 0 and 31999), got 49406. This may result in unexpected behavior.
[transformers] Model config: eos_token_id must be `None` or an integer within the vocabulary (between 0 and 31999), got 49407. This may result in unexpected behavior.
Loading weights: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 792/792 [00:01<00:00, 629.51it/s]
logit_scale=107.9853, logit_bias=-16.3476
image_embeds shape: (1, 1024)
text_embeds shape: (1, 1024)
余弦相似度: 0.1648
匹配概率(sigmoid): 0.8100