#!/usr/bin/env python3
|
# -*- coding: utf-8 -*-
|
"""导出 PyTorch 权重为 ONNX，并可选 INT8 动态量化。"""
|
import argparse
|
import json
|
import os
|
import shutil
|
import subprocess
|
import sys
|
from datetime import datetime
|
from pathlib import Path
|
|
import torch
|
import yaml
|
from train import build_model
|
|
|
def load_config(path):
|
with open(path, encoding="utf-8") as f:
|
cfg = yaml.safe_load(f)
|
base = Path(path).resolve().parent
|
out = cfg["export"]["output_dir"]
|
if not os.path.isabs(out):
|
cfg["export"]["output_dir"] = str((base / out).resolve())
|
return cfg
|
|
|
def try_quantize_subprocess(onnx_path: str, int8_path: str) -> bool:
|
"""在子进程执行 ORT 量化，避免主进程因 Windows 上 ORT bug 崩溃。"""
|
code = (
|
"import sys\n"
|
"from onnxruntime.quantization import QuantType, quantize_dynamic\n"
|
"quantize_dynamic(sys.argv[1], sys.argv[2], weight_type=QuantType.QUInt8)\n"
|
"print('OK')\n"
|
)
|
try:
|
result = subprocess.run(
|
[sys.executable, "-c", code, onnx_path, int8_path],
|
capture_output=True,
|
text=True,
|
timeout=600,
|
)
|
except subprocess.TimeoutExpired:
|
print(f"量化超时: {int8_path}")
|
return False
|
if result.returncode != 0:
|
err = (result.stderr or result.stdout or "").strip()
|
if err:
|
print(f"量化失败 exit={result.returncode}: {err[:500]}")
|
else:
|
print(f"量化失败 exit={result.returncode}（Windows 上 ORT 量化器可能崩溃）")
|
return False
|
return os.path.isfile(int8_path) and os.path.getsize(int8_path) > 0
|
|
|
def export_one(task, cfg, do_quantize: bool):
|
out_dir = cfg["export"]["output_dir"]
|
os.makedirs(out_dir, exist_ok=True)
|
size = cfg["model"]["image_size"]
|
ckpt = os.path.join(out_dir, f"{task}.pt")
|
if not os.path.isfile(ckpt):
|
raise FileNotFoundError(f"未找到权重 {ckpt}，请先运行 train.py")
|
|
model = build_model()
|
state = torch.load(ckpt, map_location="cpu")
|
model.load_state_dict(state)
|
model.eval()
|
|
dummy = torch.randn(1, 3, size, size)
|
onnx_path = os.path.join(out_dir, f"{task}.onnx")
|
torch.onnx.export(
|
model,
|
dummy,
|
onnx_path,
|
input_names=["input"],
|
output_names=["logits"],
|
dynamic_axes={"input": {0: "batch"}, "logits": {0: "batch"}},
|
opset_version=cfg["export"]["opset"],
|
)
|
print(f"导出 float ONNX: {onnx_path}")
|
|
int8_path = os.path.join(out_dir, f"{task}_int8.onnx")
|
if do_quantize and cfg["export"].get("quantize"):
|
if try_quantize_subprocess(onnx_path, int8_path):
|
print(f"量化 INT8: {int8_path}")
|
return os.path.basename(int8_path)
|
print(
|
f"WARN: {task}_int8.onnx 量化未成功，推理服务将使用 float 模型 {task}.onnx\n"
|
" 常见原因: Windows 上 onnxruntime.quantization 崩溃。"
|
" 可升级/降级 onnxruntime，或使用 --no-quantize 跳过量化。"
|
)
|
elif do_quantize:
|
print(f"跳过量化（config quantize=false）")
|
|
return os.path.basename(onnx_path)
|
|
|
def main():
|
parser = argparse.ArgumentParser(description="导出 storefront/handover ONNX 模型")
|
parser.add_argument("-c", "--config", default=str(Path(__file__).parent / "config.yaml"))
|
parser.add_argument("--version", default="1.0.0")
|
parser.add_argument("--no-quantize", action="store_true", help="仅导出 float .onnx，不尝试 INT8")
|
args = parser.parse_args()
|
cfg = load_config(args.config)
|
do_quantize = not args.no_quantize
|
|
paths = {}
|
for task in ("storefront", "handover"):
|
paths[task] = export_one(task, cfg, do_quantize)
|
|
version_info = {
|
"model_version": args.version,
|
"storefront_model": paths["storefront"],
|
"handover_model": paths["handover"],
|
"image_size": cfg["model"]["image_size"],
|
"quantized": do_quantize and paths["storefront"].endswith("_int8.onnx"),
|
"exported_at": datetime.utcnow().isoformat() + "Z",
|
}
|
version_path = os.path.join(cfg["export"]["output_dir"], "version.json")
|
with open(version_path, "w", encoding="utf-8") as f:
|
json.dump(version_info, f, indent=2, ensure_ascii=False)
|
print(f"写入 {version_path}")
|
print(f" storefront -> {paths['storefront']}")
|
print(f" handover -> {paths['handover']}")
|
|
|
if __name__ == "__main__":
|
main()
|
