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

我不是女神ヾ 2022-11-14 13:27 334阅读 0赞

Pytorch模型转换到onnx模型代码如下：

import torch
    import torch.nn as nn
    import torch.onnx
    import onnx
    import os
    from QualityNet import QualityNet
    
    if __name__ == '__main__':
    
        os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
        os.environ["CUDA_VISIBLE_DEVICES"] = "1"
    
        device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
        model_path = './models/pytorch/face_quality.pth'
        
        state_dict = torch.load(model_path)
        
        model = QualityNet()
        model = nn.DataParallel(model)
        model = QualityNet().to(device)
        model.load_state_dict(state_dict)
        model.eval()
    
        onnx_path = './models/onnx/face_quality.onnx'
        dummy_input = torch.ones(1, 3, 128, 128,)
        dummy_input = dummy_input.to(device)
        input_names = ["input"]
        output_names = ["output"]
        # export onnx model
        torch.onnx.export(model, dummy_input, onnx_path, verbose=False, opset_version=9, input_names=input_names, output_names=output_names)
        # load onnx model
        onnx_model = onnx.load(onnx_path)
        # check onnx model
        onnx.checker.check_model(onnx_model)

运行后会出现如下错误：

![20210329181841773.png][]

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

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dlaWNhbzE5OTA_size_16_color_FFFFFF_t_70][]

解决方法很简单，只需要去掉多余的"module."字段即可，重新创建一个OrderedDict，修改模型键值，然后将它载入模型，修改后的代码如下：

import torch
    import torch.nn as nn
    from collections import OrderedDict
    import torch.onnx
    import onnx
    import os
    from QualityNet import QualityNet
    
    if __name__ == '__main__':
    
        os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
        os.environ["CUDA_VISIBLE_DEVICES"] = "1"
    
        device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
        model_path = './models/pytorch/face_quality.pth'
        
        state_dict = torch.load(model_path)
        new_state_dict = OrderedDict()
        for k, v in state_dict.items():
            name = k[7:]  # remove "module."
            new_state_dict[name] = v
        
        model = QualityNet()
        model = QualityNet().to(device)
        model.load_state_dict(new_state_dict)
        model.eval()
    
        onnx_path = './models/onnx/face_quality.onnx'
        dummy_input = torch.ones(1, 3, 128, 128,)
        dummy_input = dummy_input.to(device)
        input_names = ["input"]
        output_names = ["output"]
        # export onnx model
        torch.onnx.export(model, dummy_input, onnx_path, verbose=False, opset_version=9, input_names=input_names, output_names=output_names)
        # load onnx model
        onnx_model = onnx.load(onnx_path)
        # check onnx model
        onnx.checker.check_model(onnx_model)

[20210329181841773.png]: /images/20221022/1e969dcfef2c4e86bbc29dd01bd40178.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dlaWNhbzE5OTA_size_16_color_FFFFFF_t_70]: /images/20221022/df7b93783bef4bae849b7c5060eca5af.png