Pytorch采用DataParallel进行多卡训练得到的模型文件直接转换到onnx模型会出现不支持的情况

Godlovesea

于 2024-01-21 16:51:26 发布

阅读量542

点赞数 7

文章标签：算法

本文链接：https://blog.youkuaiyun.com/weixin_42064949/article/details/135731581

版权

本文介绍了在Pytorch中使用DataParallel进行多卡训练后，如何正确地将模型转换为ONNX模型，特别提到了由于DataParallel结构导致的model.key变为models.key，并给出了相应的解决方法：在torch.onnx.export时，将model参数改为model.module。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

Pytorch DataParallel多卡训练模型导出onnx模型

Pytorch采用DataParallel进行多卡训练得到的模型文件直接转换到onnx模型会出现不支持的情况，原因是使用DataParallel进行多卡训练，模型文件中的键值对key值前面会多一个"modules.":

def export_onnx(model, im, file, opset=12):
    f = file.with_suffix('.onnx')
    torch.onnx.export(model.module, im, f, verbose=False, opset_version=opset,
    do_constant_folding=True,
    input_names=["images"],
    output_names=["da", "ll"])
    model_onnx = onnx.load(f)  # load onnx model
    onnx.checker.check_model(model_onnx)  # check onnx model
    onnx.save(model_onnx, f)

将torch.onnx.export(model 。。。。）第一个参数修改为model.module