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

将 PyTorch 训练出的 PySlowFast.pyth模型转换为 ONNX 格式的关键在于模型是双路径Dual-Pathway输入。因此在导出 ONNX 时我们需要构建一个包含两个 TensorSlow 路径和 Fast 路径的元组作为 dummy input虚拟输入。以下是完整的转换代码脚本importtorchimportnumpyasnp# 1. 定义与训练时完全一致的模型结构以标准的 SlowFast R50 为例defload_slowfast_model(checkpoint_path,num_classes10):# 使用 pytorchvideo 快速构建符合 SlowFast 架构的网络modeltorch.hub.load(facebookresearch/pytorchvideo,slowfast_r50,pretrainedFalse)# 修改分类头model.blocks[6].projtorch.nn.Linear(model.blocks[6].proj.in_features,num_classes)# 加载权重checkpointtorch.load(checkpoint_path,map_locationcpu)ifmodel_stateincheckpoint:model.load_state_dict(checkpoint[model_state])else:model.load_state_dict(checkpoint)model.eval()returnmodel# 2. 初始化模型CHECKPOINT_PATHcheckpoints/checkpoint_epoch_00050.pythNUM_CLASSES10# 替换为你的分类数modelload_slowfast_model(CHECKPOINT_PATH,num_classesNUM_CLASSES)# 3. 构造虚拟输入 (Dummy Input)# 假设输入视频为 32 帧采样率/结构对应# Slow 路径: 4 帧, 尺寸 256x256 (BatchSize1, 通道3)# Fast 路径: 32 帧, 尺寸 256x256 (BatchSize1, 通道3)batch_size1channels3height,width256,256dummy_slow_inputtorch.randn(batch_size,channels,4,height,width)dummy_fast_inputtorch.randn(batch_size,channels,32,height,width)# SlowFast 期望的输入是一个包含两个路经 Tensor 的列表或元组dummy_inputs(dummy_slow_input,dummy_fast_input)# 4. 执行 ONNX 导出onnx_output_pathslowfast_model.onnxprint(f正在转换模型并保存至:{onnx_output_path}...)torch.onnx.export(model,argsdummy_inputs,# 传入双路径虚拟输入fonnx_output_path,# 输出的 ONNX 文件路径export_paramsTrue,# 导出训练后的权重参数opset_version12,# 建议使用 opset 11 或 12对视频 3D 算子支持较好do_constant_foldingTrue,# 是否执行常量折叠优化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_modelonnx.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_sessionort.InferenceSession(slowfast_model.onnx)# 准备好你的两个经过预处理的 numpy 数组outputsort_session.run([output],{slow_input:slow_tensor_np,fast_input:fast_tensor_np})