当前位置：首页 > news >正文

将 PyTorch 训练出的 PySlowFast（.pyth）模型转换为 ONNX 格式

news 2026/7/11 1:31:52

将 PyTorch 训练出的 PySlowFast（.pyth）模型转换为 ONNX 格式的关键在于：模型是双路径（Dual-Pathway）输入。因此，在导出 ONNX 时，我们需要构建一个包含两个 Tensor（Slow 路径和 Fast 路径）的元组作为 dummy input（虚拟输入）。

以下是完整的转换代码脚本：

importtorchimportnumpyasnp# 1. 定义与训练时完全一致的模型结构（以标准的 SlowFast R50 为例）defload_slowfast_model(checkpoint_path,num_classes=10):# 使用 pytorchvideo 快速构建符合 SlowFast 架构的网络model=torch.hub.load('facebookresearch/pytorchvideo','slowfast_r50',pretrained=False)# 修改分类头model.blocks[6].proj=torch.nn.Linear(model.blocks[6].proj.in_features,num_classes)# 加载权重checkpoint=torch.load(checkpoint_path,map_location="cpu")if"model_state"incheckpoint:model.load_state_dict(checkpoint["model_state"])else:model.load_state_dict(checkpoint)model.eval()returnmodel# 2. 初始化模型CHECKPOINT_PATH="checkpoints/checkpoint_epoch_00050.pyth"NUM_CLASSES=10# 替换为你的分类数model=load_slowfast_model(CHECKPOINT_PATH,num_classes=NUM_CLASSES)# 3. 构造虚拟输入 (Dummy Input)# 假设输入视频为 32 帧，采样率/结构对应：# Slow 路径: 4 帧, 尺寸 256x256 (BatchSize=1, 通道=3)# Fast 路径: 32 帧, 尺寸 256x256 (BatchSize=1, 通道=3)batch_size=1channels=3height,width=256,256dummy_slow_input=torch.randn(batch_size,channels,4,height,width)dummy_fast_input=torch.randn(batch_size,channels,32,height,width)# SlowFast 期望的输入是一个包含两个路经 Tensor 的列表或元组dummy_inputs=(dummy_slow_input,dummy_fast_input)# 4. 执行 ONNX 导出onnx_output_path="slowfast_model.onnx"print(f"正在转换模型并保存至:{onnx_output_path}...")torch.onnx.export(model,args=dummy_inputs,# 传入双路径虚拟输入f=onnx_output_path,# 输出的 ONNX 文件路径export_params=True,# 导出训练后的权重参数opset_version=12,# 建议使用 opset 11 或 12，对视频 3D 算子支持较好do_constant_folding=True,# 是否执行常量折叠优化input_names=['slow_input','fast_input'],# 指定输入节点的名称output_names=['output'],# 指定输出节点的名称dynamic_axes={# 支持动态 BatchSize（可选）'slow_input':{0:'batch_size'},'fast_input':{0:'batch_size'},'output':{0:'batch_size'}})print("ONNX 模型转换成功！")# 5. 验证 ONNX 模型有效性importonnx onnx_model=onnx.load(onnx_output_path)onnx.checker.check_model(onnx_model)print("ONNX 模型结构检查完毕，未发现异常。")

💡 关键踩坑与避雷指南

3D 算子不兼容问题：
SlowFast 内部大量使用了 3D 卷积（Conv3d）和 3D 池化（MaxPool3d）。部分旧版本的 ONNX Opset 对其支持不佳。如果导出报错，请尝试将opset_version提升至12、13或14。
PyTorchVideo 依赖兼容性：
如果您在导出时遇到类似TracerWarning导致导出失败，可以尝试使用官方的原版 PySlowFast 代码库中的slowfast/models/build.py来构建model，而非pytorchvideo。
推理时如何传参（以 ONNX Runtime 为例）：
在后续部署时，由于设置了两个输入节点，你需要将处理好的 Slow 路径和 Fast 路径 Tensor 分别作为字典的 Key 传给推理器：

importonnxruntimeasort ort_session=ort.InferenceSession("slowfast_model.onnx")# 准备好你的两个经过预处理的 numpy 数组outputs=ort_session.run(["output"],{"slow_input":slow_tensor_np,"fast_input":fast_tensor_np})