Files
HomeAgent/scripts/export_qwen3vl_embedding_onnx.py
JianFeeeee 1a02971f88 feat(memory): 千问三段式 ONNX 嵌入补齐——可复现导出脚本 + Go 侧首次完整验证
此前三段式拆分后 ONNX 路径从未从 Go 侧跑通:embedder_onnx_test.go 仍引用
分段前的 API(e.renderInput、TextTower.onnx、旧目录),go vet -tags onnxruntime
直接编译失败。导出脚本只在 /tmp 且硬编码本机路径、从第三个目录拷贝固定形状的
Vision.onnx,完全不可复现。音频会被视觉塔编码,静默往统一空间灌入错误坐标。

本提交补齐这些缺口:

一、可复现导出脚本(scripts/export_qwen3vl_embedding_onnx.py)
- 自动拉取模型(HuggingFace 优先,失败回落 ModelScope,支持 HF_ENDPOINT 镜像);
- 导出 TokenEmbedding + Transformer + Vision 三段图,图文共用同一 token
  embedding、28 层 Transformer、last-token 池化与 fingerprint;
- 双重自检(不可省):分段 PyTorch vs 完整模型 + 导出后的 ONNX vs 完整模型,
  cos < 0.999999 即非零退出——「能加载」不等于「算得对」;
- 默认把 L2 归一化后的冻结参考向量写入产物目录(qwen_reference.json)——
  Go 测试据此做逐维冻结回归,且「该目录是哪次导出的」从文件本身可追溯;
- --verify-only 校验既有产物不重新导出,可用来确认线上在用的图没坏。

关键实测结论(已写入 docs/zh/multimodal-space.md 与长期记忆):
原生多帧视频不可行——Qwen3-VL 视觉塔把 grid_thw 当 Python 值消费
(grid_thw.tolist()),legacy tracer 固化为常量,导出后图中根本没有 grid_thw
输入,换帧数调用直接 Invalid input name: grid_thw。故视觉塔固定 (1,48,48),
视频由上层抽帧后逐帧按图像编码(同模型/同维度/同 fingerprint),音频明确
unsupported。

二、模态边界(vector.ErrModalityUnsupported)
- 新增 vector.ErrModalityUnsupported:表示「该模态不在本统一空间的原生覆盖
  范围内」,与普通错误语义不同——调用方应把它当「永远不会有向量」而非
  「本次失败、下次重试」;
- qwen.EmbedImageDense 按 mime 拒绝 audio/* 与 video/*:此前它会拿视觉塔
  去解音频字节,往统一空间灌入语义错误的坐标且静默;
- reembedStaleMedia 对 ErrModalityUnsupported 不计失败、不重试、不用别的
  模型向量顶替(TestReembedStaleMedia_SkipsUnsupportedWithoutFaking 守住)。

三、Go ONNX 测试首次完整通过
- 重写 embedder_onnx_test.go:修复编译 + 文本冻结回归 + 图像冻结回归 +
  两条阴性对照(不同输入必须不同、图像与文本必须不同)+ 不支持模态断言;
- 参考值从产物目录的 qwen_reference.json 读取(不在测试里硬编码浮点);
- 用线上部署产物实测全部通过(text cos=0.999999940, image cos=0.999999762)。

四、.gitignore 修复
- /scripts/ 此前被列在「运行时产物」下,但它是作者维护的工具目录
  (模型导出、侧车、部署校验),deploy/systemd/embed-sidecar.service 直接
  引用 scripts/embed_sidecar.py,忽略它会让那份 unit 在别人的机器上指向
  不存在的文件。改为只忽略 __pycache__。

五、文档(docs/zh/multimodal-space.md)
- 获取/启用/产物契约/模态边界/验证/资源成本/与现有部署产物的等价性。

验证:go build ./...、go vet ./...、go vet -tags onnxruntime ./...、
go test -short 全部通过;ONNX 标签测试对线上部署产物全部通过。
2026-09-11 13:45:25 +08:00

517 lines
24 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#!/usr/bin/env python3
"""自动拉取 Qwen3-VL-Embedding-2B 并导出 HomeAgent 用的三段式 ONNX 统一向量空间。
产物(写入 --out 目录,约 8GB::
TokenEmbedding.onnx input_ids -> hidden
Transformer.onnx hidden + deepstack + RoPE + causal mask -> embedding
Vision.onnx(+.data) pixel_values -> 3 层 DeepStack + 主视觉特征
tokenizer.json / tokenizer_config.json / chat_template.jinja
embed_config.json Go 侧读取的布局与契约常量
为什么是「三段」而不是一张图
--------------------------
文本与图像共用同一 token embedding、同一 28 层 Transformer、同一 last-token
池化与同一 fingerprint分段只是部署形式。把 RoPE 与视觉特征散射留在 Go 计算,
是为了避开旧式 tracer 把 seq=598 / visual=576 烘焙进图里——那样签名上写着
dynamic_axes实际却只能用导出的那个长度运行。
Vision 为什么固定 768×768不支持原生多帧视频
--------------------------------------------
Qwen3-VL 的视觉塔把 ``grid_thw`` 当 Python 值消费(``grid_thw.tolist()``
legacy tracer 会把它的内容固化成常量:实测导出后 ONNX 图里根本没有
``grid_thw`` 输入,用别的帧数调用会直接报 Invalid input name。因此这里把
grid 固定为 (1, 48, 48),并在导出处做 PyTorch↔ONNX 一致性校验。
视频由上层抽帧后逐帧按图像编码——同一模型、同一维度、同一 fingerprint
只是不做跨帧时序注意力;音频不在本空间覆盖范围内(见 vector.ErrModalityUnsupported
自检是不可省的
------------
导出脚本必须自己证明产物正确,而不是只比较有没有报错:
1. 分段 PyTorchTokenEmbedding+Transformer+Vision 的组合)对比完整模型前向;
2. 再用 onnxruntime 跑导出后的三段图,对比完整模型前向。
两步都要求 cos ≥ 0.999999,否则以非零码退出——绝不产出一个「能加载但算错」的模型。
"""
from __future__ import annotations
import argparse
import json
import os
import shutil
import sys
import time
import numpy as np
import torch
INSTRUCTION = "Represent the user's input."
IMAGE_SIZE = 768
PATCH_SIZE = 16
TEMPORAL_PATCH = 2
SPATIAL_MERGE = 2
MAX_LENGTH = 1024 # 768x768 有 576 个视觉 token512 会截断视觉占位符
DEFAULT_MODEL_ID = "Qwen/Qwen3-VL-Embedding-2B"
REFERENCE_TEXT = "今天天气怎么样"
REFERENCE_IMAGE_RGB = (200, 30, 30)
def log(msg: str) -> None:
print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)
# ─────────────────────────── 模型获取 ───────────────────────────
def pull_model(model_id: str, store: str) -> str:
"""把模型拉到本地并返回快照目录。优先 HuggingFace失败回落 ModelScope。
两个源都必须显式给出目标目录:默认的 HF blobs+snapshots 结构会把权重存成
符号链接树,导出脚本按固定文件名读取时很容易踩到不存在的路径。
HF_ENDPOINT 会被 huggingface_hub 自动识别,因此国内镜像(如
https://hf-mirror.com )无需额外参数。
"""
os.makedirs(store, exist_ok=True)
local = os.path.join(store, model_id.replace("/", "--"))
marker = os.path.join(local, "config.json")
if os.path.exists(marker):
log(f"复用已下载模型: {local}")
return local
errors = []
try:
from huggingface_hub import snapshot_download
log(f"从 HuggingFace 拉取 {model_id} -> {local}HF_ENDPOINT={os.environ.get('HF_ENDPOINT', '默认')}")
snapshot_download(repo_id=model_id, local_dir=local, max_workers=4)
if os.path.exists(marker):
return local
errors.append("huggingface: 下载完成但缺少 config.json")
except Exception as exc: # noqa: BLE001 - 需要回落到 ModelScope
errors.append(f"huggingface: {exc}")
log(f"HuggingFace 拉取失败:{exc}")
try:
from modelscope import snapshot_download as ms_snapshot
log(f"从 ModelScope 拉取 {model_id} -> {local}")
ms_snapshot(model_id, local_dir=local)
if os.path.exists(marker):
return local
errors.append("modelscope: 下载完成但缺少 config.json")
except Exception as exc: # noqa: BLE001
errors.append(f"modelscope: {exc}")
log(f"ModelScope 拉取失败:{exc}")
raise SystemExit("模型拉取失败:\n - " + "\n - ".join(errors))
# ─────────────────────────── 模型分段 ───────────────────────────
class TokenEmbedding(torch.nn.Module):
def __init__(self, embed_tokens):
super().__init__()
self.embed_tokens = embed_tokens
def forward(self, input_ids):
return self.embed_tokens(input_ids)
class Transformer(torch.nn.Module):
"""28 层语言模型 + 前置 3 层 DeepStack 相加 + final norm + last-token 池化。
池化放在图里(而不是 Go是有意的last-token 的位置由 attention_mask 决定,
一旦 Go 侧算错位置就会静默取到 padding 的 hidden而向量照样归一化、照样
能比余弦——那种错误只能靠与参考向量对比才能发现。
"""
def __init__(self, lm):
super().__init__()
self.layers = lm.layers
self.norm = lm.norm
def forward(self, hidden, deepstack_0, deepstack_1, deepstack_2,
rotary_cos, rotary_sin, causal_mask):
deep = (deepstack_0, deepstack_1, deepstack_2)
for i, layer in enumerate(self.layers):
hidden = layer(
hidden_states=hidden,
attention_mask=causal_mask,
position_embeddings=(rotary_cos, rotary_sin),
use_cache=False,
)
if i < 3:
hidden = hidden + deep[i]
return self.norm(hidden)[:, -1]
class VisionTower(torch.nn.Module):
"""视觉塔:固定 grid 的 patch 张量 -> 主视觉特征 + 3 层 DeepStack。
grid_thw 作为 buffer 固化。原因见模块 docstringlegacy tracer 无法把
grid_thw 保留为运行时输入,写成输入只会得到一个实际不含该输入的图。
"""
def __init__(self, visual, grid_thw):
super().__init__()
self.visual = visual
self.register_buffer("grid_thw", grid_thw)
def forward(self, pixel_values):
out = self.visual(pixel_values, grid_thw=self.grid_thw, return_dict=True)
d = out.deepstack_features
return d[0], d[1], d[2], out.pooler_output
# ─────────────────────────── 输入构造 ───────────────────────────
def render_text(instruction: str, text: str) -> str:
"""与 Go 侧 tokenizer.renderInstructionInput 逐字符一致。
指令放 system、正文放 user、以 assistant 起始符结尾。差一个特殊 token
last-token 池化取到的位置就变了,嵌入也就不同——而且不会报错。
"""
return (f"<|im_start|>system\n{instruction}<|im_end|>\n"
f"<|im_start|>user\n{text}<|im_end|>\n<|im_start|>assistant\n")
def text_position_ids(attention_mask):
"""无 padding 的单批语义下,三个 RoPE 轴都等于累计可见位置。"""
pos = attention_mask.long().cumsum(-1) - 1
return pos.clamp(min=0).unsqueeze(0).expand(3, -1, -1)
def text_inputs(processor, lm, text):
rendered = render_text(INSTRUCTION, text)
x = processor.tokenizer([rendered], return_tensors="pt", truncation=True,
max_length=MAX_LENGTH, padding=True)
pos = text_position_ids(x["attention_mask"])
hidden = lm.embed_tokens(x["input_ids"])
cos, sin = lm.rotary_emb(hidden, pos)
zero = torch.zeros_like(hidden)
return x, hidden, (zero, zero, zero), cos, sin, causal_mask(hidden.shape[1]), pos
def image_inputs(processor, model, image):
conv = [{"role": "system", "content": [{"type": "text", "text": INSTRUCTION}]},
{"role": "user", "content": [{"type": "image", "image": image}]}]
rendered = processor.apply_chat_template([conv], add_generation_prompt=True, tokenize=False)
x = processor(text=rendered, images=[image], do_resize=False,
return_tensors="pt", truncation=True, max_length=MAX_LENGTH)
lm = model.model.language_model
with torch.no_grad():
vo = model.model.visual(x["pixel_values"], grid_thw=x["image_grid_thw"], return_dict=True)
pos, _ = model.model.get_rope_index(
x["input_ids"], x["mm_token_type_ids"],
image_grid_thw=x["image_grid_thw"], attention_mask=x["attention_mask"])
mask = x["mm_token_type_ids"] == 1
hidden = lm.embed_tokens(x["input_ids"])
hidden = hidden.clone()
hidden[mask] = vo.pooler_output
deep = []
for d in vo.deepstack_features:
full = torch.zeros_like(hidden)
full[mask] = d
deep.append(full)
cos, sin = lm.rotary_emb(hidden, pos)
return x, hidden, tuple(deep), cos, sin, causal_mask(hidden.shape[1]), pos
def causal_mask(seq: int) -> torch.Tensor:
m = torch.full((1, 1, seq, seq), torch.finfo(torch.float32).min)
return torch.triu(m, diagonal=1)
def pool_last(hidden, mask):
last = mask.shape[1] - mask.flip(1).argmax(1) - 1
return hidden[torch.arange(hidden.shape[0], device=hidden.device), last]
# ─────────────────────────── 校验 ───────────────────────────
def compare(name: str, a, b, floor: float = 0.999999) -> float:
a = torch.nn.functional.normalize(torch.as_tensor(a).float(), dim=-1)
b = torch.nn.functional.normalize(torch.as_tensor(b).float(), dim=-1)
cos = float((a * b).sum())
md = float((a - b).abs().max())
log(f" {name}: cos={cos:.9f} maxdiff={md:.3e}")
if cos < floor:
raise SystemExit(f"导出校验失败:{name} 与完整模型不等价 (cos={cos:.9f})")
return cos
def verify_split_torch(model, processor, transformer) -> None:
log("校验①:分段 PyTorch vs 完整模型")
x, h, d, cos, sin, cm, pos = text_inputs(processor, model.model.language_model, REFERENCE_TEXT)
with torch.no_grad():
got = transformer(h, *d, cos, sin, cm)
ref = model.model(input_ids=x["input_ids"], attention_mask=x["attention_mask"],
position_ids=pos, use_cache=False).last_hidden_state[:, -1]
compare("text/split-torch", got, ref)
img = reference_image()
x, h, d, cos, sin, cm, _ = image_inputs(processor, model, img)
with torch.no_grad():
got = transformer(h, *d, cos, sin, cm)
ref = model.model(input_ids=x["input_ids"], attention_mask=x["attention_mask"],
pixel_values=x["pixel_values"], image_grid_thw=x["image_grid_thw"],
mm_token_type_ids=x["mm_token_type_ids"],
use_cache=False).last_hidden_state[:, -1]
compare("image/split-torch", got, ref)
log(f" 形状: seq={h.shape[1]} visual={(x['mm_token_type_ids'] == 1).sum().item()}")
def unit(vec):
"""L2 归一化。
必须对**写出的参考向量**归一化ONNX 图返回的是 final norm 之后的原始 last hidden
而 Go 侧的 public 接口返回的是归一化后的向量。若参考用原始值Go 测试会全线不匹配——
且这个差异看起来像“模型不对”,实际上只是两边对“向量”的定义不同。
"""
v = np.asarray(vec, dtype=np.float64)
n = float(np.linalg.norm(v))
return v / n if n > 0 else v
def verify_onnx(out_dir: str, model, processor) -> dict:
"""用 onnxruntime 跑导出后的三段图,对比完整模型前向。
返回参考向量(供 Go 侧测试冻结使用Go 必须复现同一套预处理与模板,
因此这里把同一输入下的期望向量前若干维导出。
"""
import onnxruntime as ort
log("校验②:导出后的 ONNX 三段图 vs 完整模型")
ts = ort.InferenceSession(os.path.join(out_dir, "TokenEmbedding.onnx"), providers=["CPUExecutionProvider"])
xs = ort.InferenceSession(os.path.join(out_dir, "Transformer.onnx"), providers=["CPUExecutionProvider"])
vs = ort.InferenceSession(os.path.join(out_dir, "Vision.onnx"), providers=["CPUExecutionProvider"])
lm = model.model.language_model
reference: dict[str, object] = {}
def run_transform(hidden, deep, cos, sin):
seq = hidden.shape[1]
return xs.run(None, {
"hidden": hidden.astype(np.float32),
"deepstack_0": deep[0].astype(np.float32),
"deepstack_1": deep[1].astype(np.float32),
"deepstack_2": deep[2].astype(np.float32),
"rotary_cos": cos.astype(np.float32),
"rotary_sin": sin.astype(np.float32),
"causal_mask": causal_mask(seq).numpy(),
})[0]
x, h, d, cos, sin, _cm, pos = text_inputs(processor, lm, REFERENCE_TEXT)
with torch.no_grad():
ref_text = model.model(input_ids=x["input_ids"], attention_mask=x["attention_mask"],
position_ids=pos, use_cache=False).last_hidden_state[:, -1].numpy()
h_onnx = ts.run(None, {"input_ids": x["input_ids"].numpy().astype(np.int64)})[0]
zero = np.zeros_like(h_onnx)
got = run_transform(h_onnx, (zero, zero, zero), cos.numpy(), sin.numpy())
compare("text/onnx-vs-full", got, ref_text)
reference["text"] = REFERENCE_TEXT
reference["text_vector_prefix"] = [float(v) for v in unit(got[0])[:12]]
reference["text_norm_raw"] = float(np.linalg.norm(got[0]))
img = reference_image()
x, _h, _d, cos, sin, _cm, _ = image_inputs(processor, model, img)
with torch.no_grad():
ref_img = model.model(input_ids=x["input_ids"], attention_mask=x["attention_mask"],
pixel_values=x["pixel_values"], image_grid_thw=x["image_grid_thw"],
mm_token_type_ids=x["mm_token_type_ids"],
use_cache=False).last_hidden_state[:, -1].numpy()
h_onnx = ts.run(None, {"input_ids": x["input_ids"].numpy().astype(np.int64)})[0]
vo = vs.run(None, {"pixel_values": x["pixel_values"].numpy().astype(np.float32)})
mask = x["mm_token_type_ids"].numpy() == 1
h_onnx = h_onnx.copy()
h_onnx[mask] = vo[3]
deep = []
for d in vo[:3]:
full = np.zeros_like(h_onnx)
full[mask] = d
deep.append(full)
got = run_transform(h_onnx, tuple(deep), cos.numpy(), sin.numpy())
compare("image/onnx-vs-full", got, ref_img)
reference["image_rgb"] = list(REFERENCE_IMAGE_RGB)
reference["image_size"] = IMAGE_SIZE
reference["image_vector_prefix"] = [float(v) for v in unit(got[0])[:12]]
reference["image_norm_raw"] = float(np.linalg.norm(got[0]))
reference["dim"] = int(got.shape[1])
return reference
def reference_image():
from PIL import Image
return Image.new("RGB", (IMAGE_SIZE, IMAGE_SIZE), REFERENCE_IMAGE_RGB)
# ─────────────────────────── 导出 ───────────────────────────
def export_graphs(out_dir: str, model, processor, model_dir: str, transformer) -> None:
os.makedirs(out_dir, exist_ok=True)
# 清掉旧产物,避免 fingerprint 把死文件算进去(旧图/旧外部权重会让
# 空间指纹变化,触发一次毫无意义的全量重算)。
for name in os.listdir(out_dir):
p = os.path.join(out_dir, name)
if os.path.isfile(p):
os.remove(p)
lm = model.model.language_model
log("导出 TokenEmbedding.onnx")
ids = torch.tensor([[151643, 151643]], dtype=torch.long)
with torch.no_grad():
torch.onnx.export(
TokenEmbedding(lm.embed_tokens).eval(), (ids,), os.path.join(out_dir, "TokenEmbedding.onnx"),
input_names=["input_ids"], output_names=["hidden"],
dynamic_axes={"input_ids": {1: "seq"}, "hidden": {1: "seq"}},
opset_version=17, do_constant_folding=True, dynamo=False,
)
log("导出 Transformer.onnx")
x, h, d, cos, sin, cm, _ = image_inputs(processor, model, reference_image())
with torch.no_grad():
torch.onnx.export(
transformer, (h, *d, cos, sin, cm), os.path.join(out_dir, "Transformer.onnx"),
input_names=["hidden", "deepstack_0", "deepstack_1", "deepstack_2",
"rotary_cos", "rotary_sin", "causal_mask"],
output_names=["embedding"],
dynamic_axes={
"hidden": {1: "seq"}, "deepstack_0": {1: "seq"}, "deepstack_1": {1: "seq"},
"deepstack_2": {1: "seq"}, "rotary_cos": {1: "seq"}, "rotary_sin": {1: "seq"},
"causal_mask": {2: "seq", 3: "seq"},
},
opset_version=17, do_constant_folding=True, dynamo=False,
)
log("导出 Vision.onnx固定 grid 1×48×48")
grid = torch.tensor([[1, IMAGE_SIZE // PATCH_SIZE, IMAGE_SIZE // PATCH_SIZE]], dtype=torch.long)
pv = x["pixel_values"]
with torch.no_grad():
torch.onnx.export(
VisionTower(model.model.visual, grid).eval(), (pv,), os.path.join(out_dir, "Vision.onnx"),
input_names=["pixel_values"],
output_names=["deepstack_feature_0", "deepstack_feature_1",
"deepstack_feature_2", "vision_hidden_states"],
opset_version=17, do_constant_folding=True, dynamo=False,
)
for name in ("tokenizer.json", "tokenizer_config.json", "chat_template.jinja", "added_tokens.json"):
src = os.path.join(model_dir, name)
if os.path.exists(src):
shutil.copy2(src, os.path.join(out_dir, name))
def write_config(out_dir: str, model, processor) -> None:
cfg = model.config
text_cfg = getattr(cfg, "text_config", cfg)
rope_scaling = getattr(text_cfg, "rope_scaling", None) or {}
mrope_section = rope_scaling.get("mrope_section") or [24, 20, 20]
vision = cfg.vision_config
meta = {
"arch": "qwen3-vl-embedding-2b-multimodal",
"runtime": "homeagent-onnx-three-part",
"dim": int(getattr(text_cfg, "hidden_size", 2048)),
"max_length": MAX_LENGTH,
"instruction": INSTRUCTION,
"pooling": "last_token",
"normalize": True,
"image_size": IMAGE_SIZE,
"patch_size": int(vision.patch_size),
"temporal_patch_size": int(vision.temporal_patch_size),
"spatial_merge_size": int(vision.spatial_merge_size),
"image_mean": [0.5, 0.5, 0.5],
"image_std": [0.5, 0.5, 0.5],
"rope_theta": float(getattr(text_cfg, "rope_theta", 5000000)),
"mrope_section": [int(v) for v in mrope_section],
"num_layers": int(getattr(text_cfg, "num_hidden_layers", 28)),
"supports_native_video": False,
"unsupported_modalities": ["audio", "video"],
"notes": ("视频由上层抽帧后逐帧按图像编码(同模型/同维度/同 fingerprint"
"音频需未来接入真正的统一音频模型。grid_thw 被 legacy tracer 固化为常量,"
"故视觉塔固定 1×48×48详见导出脚本 docstring。"),
}
with open(os.path.join(out_dir, "embed_config.json"), "w") as f:
json.dump(meta, f, ensure_ascii=False, indent=2)
log(f"写出 embed_config.json dim={meta['dim']} rope_theta={meta['rope_theta']} "
f"mrope={meta['mrope_section']}")
def main() -> int:
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("--out", required=True, help="ONNX 产物目录(写入约 8GB")
ap.add_argument("--model-dir", default="", help="已下载的模型目录;给出则跳过自动拉取")
ap.add_argument("--model-id", default=DEFAULT_MODEL_ID, help=f"模型仓库 id默认 {DEFAULT_MODEL_ID}")
ap.add_argument("--model-store", default=os.path.expanduser("~/.cache/homeagent-qwen-models"),
help="自动拉取时的模型存放目录")
ap.add_argument("--no-reference", action="store_true",
help="不写 <out>/qwen_reference.json默认会写Go 测试靠它做冻结回归)")
ap.add_argument("--skip-verify", action="store_true", help="跳过导出后校验(仅调试用,不推荐)")
ap.add_argument("--verify-only", action="store_true",
help="不重新导出,只校验已存在的 <out> 并(重新)写出参考向量")
args = ap.parse_args()
if args.verify_only:
# 校验既有产物目录:既能确认线上在用的图没坏,也能给旧目录补上参考向量。
for name in ("TokenEmbedding.onnx", "Transformer.onnx", "Vision.onnx"):
if not os.path.exists(os.path.join(args.out, name)):
raise SystemExit(f"{args.out} 下缺少 {name},无法 --verify-only")
model_dir = args.model_dir or pull_model(args.model_id, args.model_store)
from transformers import AutoProcessor
from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLForConditionalGeneration
log("加载 processor / model--verify-only")
processor = AutoProcessor.from_pretrained(model_dir, trust_remote_code=True, padding_side="right")
model = Qwen3VLForConditionalGeneration.from_pretrained(
model_dir, dtype=torch.float32, low_cpu_mem_usage=True).eval()
reference = verify_onnx(args.out, model, processor)
if not args.no_reference:
path = os.path.join(args.out, "qwen_reference.json")
with open(path, "w") as f:
json.dump(reference, f, ensure_ascii=False, indent=2)
log(f"写出冻结参考向量: {path}")
log(f"校验完成: {args.out}")
return 0
if args.model_dir:
model_dir = args.model_dir
if not os.path.exists(os.path.join(model_dir, "config.json")):
raise SystemExit(f"--model-dir {model_dir} 下没有 config.json")
log(f"使用本地模型: {model_dir}")
else:
model_dir = pull_model(args.model_id, args.model_store)
# transformers 只在真正导出时才需要(拉取模型本身只用 huggingface_hub
from PIL import Image # noqa: F401 确保依赖存在并给出清晰报错
from transformers import AutoProcessor
from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLForConditionalGeneration
log("加载 processor / modelFP32CPU")
processor = AutoProcessor.from_pretrained(model_dir, trust_remote_code=True, padding_side="right")
model = Qwen3VLForConditionalGeneration.from_pretrained(
model_dir, dtype=torch.float32, low_cpu_mem_usage=True).eval()
transformer = Transformer(model.model.language_model).eval()
verify_split_torch(model, processor, transformer)
export_graphs(args.out, model, processor, model_dir, transformer)
write_config(args.out, model, processor)
reference = None if args.skip_verify else verify_onnx(args.out, model, processor)
# 参考写进产物目录本身:这样任何一个 ONNX 目录都自带「它应当给出什么输出」,
# Go 测试无需额外配置就能找到,也不会出现「模型换了、参考还是旧的」的错配。
if reference is not None and not args.no_reference:
path = os.path.join(args.out, "qwen_reference.json")
with open(path, "w") as f:
json.dump(reference, f, ensure_ascii=False, indent=2)
log(f"写出冻结参考向量: {path}")
total = sum(os.path.getsize(os.path.join(args.out, n))
for n in os.listdir(args.out) if os.path.isfile(os.path.join(args.out, n)))
log(f"完成: {args.out}{total / 2**30:.2f} GiB")
log("Go 侧用法: core.memory.multimodal_space.type=onnx + "
f"core.memory.multimodal_space.onnx.model_dir={args.out}")
return 0
if __name__ == "__main__":
sys.exit(main())