PyPTO-Gym Convert-Model 实验:PyTorch 到 ONNX / TorchScript / safetensors 的跨格式转换与数值验证实战
发布时间:2026/9/18 5:48:06 锦皓数字建站

PyPTO-Gym Convert-Model 实验PyTorch 到 ONNX / TorchScript / safetensors 的跨格式转换与数值验证实战【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym本文基于 pypto-gym 仓库中cannbot-skills/model/pypto-convert-model/下的 Convert-Model 实验文档references/README.md展开讲清这套实验以 PyTorch 为基准、把模型转换到 onnx / pt / safetensors 三种格式、再重载产物并与原始前向输出做逐元素比对的完整方法论。读完你可以掌握如何在 CNN / Transformer-MLP / MoE 三类模型上批量跑格式转换矩阵、如何用run.py管理实验工作区与结果归档、如何用port.py完成任意两格式之间的转换并自动输出max_abs数值安全判定以及实验对设备自动探测、ONNX 导出参数、容差策略等关键工程细节的设计依据。1. 实验设计PyTorch 是唯一事实来源文档的第一句话定义了整套实验的基准原则PyTorch is the source of truth。具体流程是把每个模型转换convert为三种目标格式之一onnx、ptTorchScript、safetensors将转换产物重新加载reload在同一份 dummy 输入上把产物前向输出与原始 PyTorch 模型的前向输出做数值比对得到max_abs等误差指标判定本次 round-trip 是否在数值上安全。目标格式即文档中列出的onnx、ptTorchScript、safetensors三者在源码中对应一个常量列表TARGET_FORMATS [onnx, pt, safetensors]见 registry.py。三种格式的语义并不相同onnx 与 pt 是可执行图而 safetensors 只是权重容器——因此 safetensors 的重跑必须借助原始模型模板把权重灌回去再前向这一差异在后文第 5 节的比对环节有专门处理。2. 模型注册表9 个模型覆盖三类架构实验的模型清单由 registry.py 中的REGISTRY字典维护共 9 个条目按类别覆盖 CNN、Transformer-MLP 与 MoE 三类idcategorysourcereposizeresnet18CNNHuggingFacemicrosoft/resnet-18~45MBmobilenet_v2CNNHuggingFacegoogle/mobilenet_v2_1.0_224~14MBefficientnet_b0CNNtimmefficientnet_b0.ra_in1k~21MBmlp_mixerTransformer-MLPtimmmixer_b16_224.goog_in21k_ft_in1k~240MBgmlpTransformer-MLPtimmgmlp_s16_224.ra3_in1k~80MBresmlpTransformer-MLPtimmresmlp_12_224.fb_in1k~60MBswitch_base_8MoEHuggingFacegoogle/switch-base-8~300MBqwen_moeMoEHuggingFaceQwen/Qwen1.5-MoE-A2.7B~14GBtoy_topk_moe / toy_soft_moe / toy_switch_moeMoElocaltiny MoE definitions1MB注册表里每个条目不只是名字与来源还携带两条驱动整个流水线的元信息可从 registry.py 确认loader字段决定用哪条加载路径——image-classificationHF 图像分类头、timm、seq2seq-lm、causal-lm或toy本地构造input字段dummy 输入规格。图像类模型统一使用{kind: image, shape: (1, 3, 224, 224)}的随机张量语言模型使用固定 prompt 文本如 switch_base_8 的translate English to German: Hello world.、qwen_moe 的Hello, the quick brown由加载阶段 tokenize 成input_idstoy MoE 则直接给出张量形状(2, 64)。qwen_moe 额外带有一个size_warning_gb: 14标记提示该模型体量很大fp32 权重约 14GB对应 loaders.py 中为 causal LM 默认改用float16加载的内存优化可用环境变量CONVERT_EXP_DTYPEfloat32|bfloat16|float16覆盖。2.1 本地 toy MoE零下载的管线验证件三个 toy MoE 定义在 toy_moe.py用随机权重构造目的不是评估精度而是低成本地压测转换管线本身尤其是 MoE 这种含数据依赖控制流的架构TopKMoEtoy_topk_moe经典 top-k 门控 MoEdim64, num_experts4, top_k2。前向中先gate_logits.topk(top_k)再逐专家用(idx e)掩码做out[mask] w[mask] * expert(x[mask])的稀疏累加——这种按索引取子集的写法正是 ONNX dynamo 导出器难以处理的典型形态SoftMoEtoy_soft_moe软路由 MoE每个专家都看到全部输入输出为weights.unsqueeze(-1) * outs的加权和控制流上相对稠密SwitchMoEtoy_switch_moeSwitch-Transformer 风格 top-1 路由argmax选专家后逐专家掩码前向。build_toy(name)统一以torch.manual_seed(42)构造模型并返回(model, sample)sample 形状(2, 64)保证每次运行输入可复现见 toy_moe.py。3. 统一加载层把 HF / timm / toy 收敛到同一接口loaders.py 的load(model_id, cache_root)是唯一入口按注册表分派到 5 条加载路径统一返回(model, sample_input, meta)三元组模型处于 eval 模式且在 CPU 上meta 携带kind、hf_dir、input_names等上下文。几个值得注意的工程决策输出适配层TensorOutputAdapterloaders.pyHF 模型的前向通常返回带logits/last_hidden_state的 dataclass导出图时会引入大量无关节点。该适配器把输出剥离为裸张量优先取logits其次last_hidden_state再退回元组首元素并透传state_dict/load_state_dict使 safetensors 存取能直接作用到内层模型。seq2seq 与 causal LM 还分别套了_Seq2SeqLogits固定forward(input_ids, decoder_input_ids)位置签名和_CausalLogits固定forward(input_ids)确保导出图的输入签名确定缓存目录约定每个模型的下载缓存位于cache_root/category/model_id由load负责mkdir(parentsTrue, exist_okTrue)同一类别的模型共享父目录timm 加载从 repo 字符串中截取/后的模型名如efficientnet_b0.ra_in1k调用timm.create_model(name, pretrainedTrue)。4. 三个转换器各格式的实现差异与导出参数converters.py 用CONVERTERS字典把格式名映射到三个转换函数每个函数签名统一为(model, sample, out_dir, meta) - Path约定失败就抛异常由 runner 捕获并记日志。4.1 safetensorssave_model绕开 weight-tied 张量to_safetensors使用safetensors.torch.save_model而非save_file(state_dict)。源码注释写明原因使用save_model可以让 weight-tied权重共享的张量不触发 safetensors 的限制见 converters.py。对存在tied_weights的 HF 模型如部分 decoder直接序列化state_dict常因重复张量而报错save_model走模型接口天然规避了该问题。4.2 TorchScripttrace 优先、script 兜底to_pt先尝试torch.jit.trace(model, args, strictFalse)trace 失败再退回torch.jit.script(model)见 converters.py。trace 路线能吞下大量动态控制流对 toy MoE 这类含for循环 掩码赋值的模型是可行路径strictFalse则避免 trace 时对输出一致性校验过严导致的失败。4.3 ONNXopset 17 dynamoFalse的关键细节to_onnx的导出参数converters.pytorch.onnx.export( model, args, str(out), input_namesnames, output_names[output], opset_version17, dynamic_axesdynamic_axes, # 输入与输出的第 0 维标记为 batch do_constant_foldingTrue, dynamoFalse, # 关键固定使用旧版 tracing 导出器 )其中dynamoFalse正是原文档 Environment notes 中强调的点源码注释补充了具体动机torch2.10 默认切换到 dynamo 导出器而 dynamo 在数据依赖控制流如 MoE 路由上会失败因此固定回 legacy tracing 导出器。同时代码对老版本torch.onnx.export不支持dynamo参数的情况做了TypeError兜底直接去掉该参数重试。dynamic_axes把输入与输出的 batch 维声明为动态使得导出图不被 dummy 输入的 batch 大小绑死。输入名归一化由_normalize(sample, meta)完成tuple 样本按meta[input_names]或arg{i}命名dict 样本按 key 命名普通张量统一叫input——这与 loaders.py 中为 seq2seq/causal 模型显式塞入input_names[input_ids, decoder_input_ids]、[input_ids]的做法是配套的。5. 数值比对diff的完整指标与容差策略比对核心在 compare.py 中。5.1 三种产物的重跑方式onnxrun_onnx用onnxruntime.InferenceSession加载按sess.get_inputs()的顺序把样本张量喂进去运行后取第一个输出展平为 numpy 数组compare.pyptrun_pt用torch.jit.load加载优先放到自动探测到的设备上设备不可用则回退 CPUsafetensorsrun_safetensors由于产物只有权重需要重新load(model_id, CACHE)一个模型模板load_model(template, path, strictFalse)灌入权重后走与参考输出相同的前向路径。5.2diff的指标体系diff(a, b, atol1e-4, rtol1e-3)返回一个信息量很大的字典compare.py误差统计max_abs、mean_abs、max_rel以及p50 / p95 / p99分位数存在非法值时分位数直接记为 inf形状与大小一致性ref_shape/out_shape/shape_match/size_match比对前先把两个数组 ravel 并截到min(ref_size, out_size)inf 语义处理_calculate_errors区分双方都是同号 inf视为一致与其余非有限值invalid置 inf 误差避免 softmax 尾部产生的 ±inf 直接污染结论top_mismatch用np.argpartition找出误差最大的 5 个展平下标及误差值方便人工定位判定字段allclose仅当大小一致、形状一致、无非法值、且np.isclose(reference, output, atol, rtol)全真时才为真runner 据此打出PASS或DIFF。bfloat16/float16张量在转 numpy 前统一先转float32_to_numpy_safenumpy 不支持这些 dtype保证比对在 float32 空间进行。5.3 设备相关的容差选择port 工具中的_tolerances_forport.py揭示了另一层工程判断比对完全在 CPU 上完成时用更紧的atol1e-4, rtol1e-3参考前向跑在 CUDA 等加速器上时放宽到atol5e-3, rtol1e-2——注释说明即使权重与输入完全相同CUDA 的 fp32 前向因 cudnn 累加噪声本身就存在 ~5e-4 量级的波动用 CPU 级紧容差会产生假 DIFF。为此port.py顶部还固定了torch.backends.cudnn.deterministic True且benchmark False尽量消除 kernel 选择噪声让 diff 反映的是转换质量而非算子调度差异。这也呼应了原文档 Environment notes 的第三点ONNX 以 float32 运行预期对 PyTorch 参考输出的 max abs diff 在 ~1e-6 到 1e-4 量级。6. 设备自动探测NPU CUDA XPU MPS CPUdevice_util.py 用一个 picker 列表实现了原文档所述的自动探测顺序NPU CUDA/ROCm XPU MPS CPUNPU 探测先看torch.npu是否存在再兜底import torch_npu见 _has_npupick_device()支持环境变量CONVERT_EXP_DEVICE强制指定npu|cuda|xpu|mps|cpu若指定的设备不可用直接抛RuntimeError并列出可用设备清单而不是静默降级结果带模块级缓存_CACHED一次进程内只探测一次。配套的onnxruntime_providers()device_util.py返回按优先级排序、且过滤到本机实际安装的 ORT 执行提供者优先级链为TensorrtExecutionProvider→CUDAExecutionProvider→ROCMExecutionProvider→MIGraphXExecutionProvider→CANNExecutionProvider昇腾 NPU→QNNExecutionProvider高通 NPU→DmlExecutionProvider→CoreMLExecutionProvider→OpenVINOExecutionProvider→XnnpackExecutionProvider→CPUExecutionProvider兜底必加。也就是说 ONNX 产物在昇腾环境会自动走 CANN 执行提供者这正是该实验面向 CANN 场景的落点。此外 compare.py 的reference_output还处理了一个容易踩坑的细节参考前向会把模型挪到加速器上但转换阶段ONNX/TorchScript 导出假设模型在 CPU 上因此跑完必须把模型恢复到原设备若加速器上 OOM统一内存主机跑大模型常见则捕获 OOM 特征串out of memory/cudnn_status_alloc/cublas_status_alloc等后回退 CPU 继续保证转换流程不中断。7. 批量实验 runnerrun.py的工作区与产物7.1 运行方式原文档给出的批量入口source .venv/bin/activate python scripts/run.py # all models in registry python scripts/run.py mobilenet_v2 toy_moe # subsetrun.py接受任意个模型 id 作为位置参数缺省跑注册表全部模型run.py。7.2 路径约定仓库内证据 vs 运行时工作区原文档 Layout 一节描述的convert_experiment/结构在实现上被拆成了提交进仓库的冻结证据与gitignore 的运行时产物两部分见 run.py 头部 docstring 与 L44-L52# 冻结证据提交到仓库 references/matrix.md # 生成的 PASS/FAIL 矩阵 references/results/id.json # 每模型结果 # 运行时产物gitignored默认 ~/.cache/pypto-convert-model # 可用环境变量 CONVERT_MODEL_WORKDIR 覆盖 $WORKDIR/models/ # HF / timm 下载缓存 $WORKDIR/outputs/model/fmt/ # 转换产物model.onnx / model.pt / model.safetensors $WORKDIR/logs/ # master 每个转换的独立日志对应 run.py 中WORKDIR取CONVERT_MODEL_WORKDIR或~/.cache/pypto-convert-modelCACHE WORKDIR/models、OUT WORKDIR/outputs而RESULTS与MATRIX固定落在技能目录的references/下即仓库中的 matrix.md 与references/results/*.json。7.3 单模型流水线process_model(model_id)的执行顺序run.pyload()加载源模型并立即跑一次reference_output()得到参考输出源模型只加载一次三种格式共用省下载省显存对TARGET_FORMATS逐个执行_convert_and_test_single_formatCONVERTERSfmt写产物 →run_target()重跑产物 →cmp.diff()与参考比对 → 状态记为PASSallclose或DIFF每个格式完成后gc.collect()并del释放模型控制大模型实验的内存水位每模型结果写references/results/id.json。状态机完整枚举run.pyPASS / DIFF / SKIP / CONVERT_FAIL(C-FAIL) / TEST_FAIL(T-FAIL) / TEST_SKIP(T-SKIP) / LOAD_FAIL(L-FAIL) / PENDING(?)。加载阶段失败会把三个格式全部标成L-FAIL转换失败与测试失败被区分开分别捕获NotImplementedError与其他异常便于矩阵里一眼看出是导不出来还是导出来但跑不对。7.4 矩阵生成与断点续跑write_matrix()在每个模型跑完一次后就会重写整张 matrix.mdrun.py矩阵表头记录自动探测到的测试设备单元格除状态标签外还会附上max_absx.xx exx。main()启动时先读取references/results/*.json里已有的部分结果partial progress accumulates支持中断后重跑时只补跑缺的模型Ctrl-C也能保留已写出的进度。8. 实际运行证据仓库中冻结的矩阵与结果 JSON仓库提交的 matrix.md 是上一轮在 NPU 设备上跑的冻结证据modelonnxptsafetensorsresnet18 … qwen_moe8 个 HF/timm 模型L-FAILL-FAILL-FAILtoy_moeC-FAILOK, max_abs0.00e00OK, max_abs0.00e00toy_soft_moeC-FAILOK, max_abs0.00e00OK, max_abs0.00e00toy_switch_moeC-FAILOK, max_abs0.00e00OK, max_abs0.00e00这组结果本身就是一次失败模式的实证HF/timm 模型在加载阶段模型下载/构建即失败三个格式统一记L-FAILtoy 模型的 ONNX 导出记C-FAIL而 pt 与 safetensors 全部OK且max_abs0。references/results/toy_moe.json 给出了具体原因onnx 项为CONVERT_FAIL错误信息是OnnxExporterError(Module onnx is not installed!)——即该轮环境缺onnx包pt 项显示convert_seconds: 0.35、产物outputs/toy_moe/pt/model.pt0.3 MB、diffmax_abs: 0.0 / allclose: true。这提醒使用者toy 模型max_abs0的完美一致部分来自 trace 出的图与原始前向在同一设备上重放而 ONNX 路径需要完整安装onnx工具链后才可比较。9. 任意格式互转 CLIport.py除批量实验外port.py 提供一个any-to-any porter用于手里已有某个格式产物、想直接搬到另一格式的场景。原文档给出的三个典型用法# pt - onnx, 无需额外信息TorchScript 自包含 python scripts/port.py model.pt out.onnx --input-shape 1,3,224,224 # safetensors - onnx, 需要一个 HF repo 提供架构 python scripts/port.py model.safetensors out.onnx --hf-repo google/mobilenet_v2_1.0_224 # onnx - pt内部使用 onnx2torch python scripts/port.py model.onnx out.pt --input-shape 1,3,224,224完整 CLI 参数port.py参数说明input/output位置参数输入、输出产物路径格式默认按扩展名推断.onnx / .pt / .pth / .safetensors--input-format/--output-format覆盖扩展名推断取值onnx、pt、safetensors--input-shape逗号分隔的 dummy 输入形状如1,3,224,224用于构造随机输入safetensors 为输入且未给形状时默认回退1,3,224,224--hf-repo输入为 safetensors 时必需用它实例化架构先试AutoModelForImageClassification失败再退回AutoModel见 _instantiate_hf_model--skip-verify跳过 round-trip 数值比对-v/--verbose打印每阶段耗时、张量统计shape/dtype/min/max/mean/std/first5、diff 分布与 top mismatch 下标内部流程_port()port.py严格分三步加载源pt 走torch.jit.loadonnx 走onnx2torch.convert变回nn.Module才能再导出safetensors 走 HF 架构实例化 load_model→ 写目标格式 → 验证 round-trip源前向输出 vs 目标前向输出按第 5.3 节策略选容差最终status为PASS或DIFF。CLI 退出码与之一致0 PASS1 DIFF2 异常可直接用于脚本化判定。值得强调的是每次 port 都会打印max_abs差值so you immediately know if the round-trip is numerically safe——把转换工具从能跑就行提升为自带验收。10. 环境依赖与适用前提requirements.txt 声明的依赖分层清晰# Core torch2.5 transformers4.46 huggingface_hub safetensors accelerate # Image classifiers (pretrained source for the experiment) timm # ONNX path onnx1.16 onnxruntime onnxscript # Optional: needed only by port.py for onnx → pt onnx2torch由此可归纳适用前提与限制设备测试/验证阶段自动探测加速器顺序 NPU CUDA/ROCm XPU MPS CPU可用CONVERT_EXP_DEVICEnpu|cuda|xpu|mps|cpu强制指定指定的设备不可用会直接报错而非降级ONNX 导出固定dynamoFalse走 legacy tracing 导出器以支持 MoE 类数据依赖控制流opset 17batch 维动态float32 精度下预期 max abs diff ~1e-6 到 1e-4大模型内存causal LM如 14GB 的 qwen_moe默认 fp16 加载统一内存主机上参考前向 OOM 时自动回退 CPU可选依赖onnx2torch仅port.py的 onnx → pt 方向需要onnx缺失会导致 ONNX 导出C-FAIL见第 8 节实证工作区运行时产物默认落在~/.cache/pypto-convert-model可用CONVERT_MODEL_WORKDIR重定向仓库内的references/matrix.md与references/results/*.json是上一轮运行的冻结快照重跑会被更新。11. 小结这套 Convert-Model 实验的核心价值在于把模型换格式从一个手工动作变成可批量、可复现、带数值验收的实验流程registry.py声明模型与输入规格loaders.py统一加载converters.py处理三种格式各自的导出陷阱weight-tied、trace/script 回退、dynamoFalsecompare.py提供含分位数、inf 语义与 top-mismatch 的完整 diffrun.py负责矩阵化跑批与断点续跑port.py则把同一套转换 验收能力暴露给单文件场景。其冻结证据matrix.md、references/results/同时记录了成功路径与各类失败模式L-FAIL/C-FAIL 的区分对排查导不出与导出来不对两类问题都有直接参考价值。【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
锦
锦皓数字建站
深耕本土企业品牌数字化升级,专注原创端正雅致商务官网,从视觉设计到稳定运维全程保驾护航。