将 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})