pytorch模型转换为onnx

东方耀

pytorch模型转换为onnx

1. 
   
2. import torch.onnx
   
3. from config import get_config_lite
   
4. from Learner import face_learner
   
5. 
   
6. 
   
7. if __name__ == '__main__':
   
8.     # input_model = "./work_space_03_lite/save/model_2020-07-18-02-41_accuracy:0.9261636148230064_step:3311920_final.pth"
   
9.     output_onnx = './work_space_03_lite/save/FaceRecog_dfy.onnx'
   
10.     print("==> Exporting model to ONNX format at '{}'".format(output_onnx))
    
11.     conf = get_config_lite(training=False)
    
12.     # conf = get_config_dfy(training=False)
    
13.     print("easydict的配置参数：", conf)
    
14.     learner = face_learner(conf, inference=True)
    
15.     fixed_str = "2020-07-18-02-41_accuracy:0.9261636148230064_step:3311920_final.pth"
    
16.     learner.load_state(conf, fixed_str, model_only=True, from_save_folder=True)
    
17.     # 开始：pytorch模型转换为onnx
    
18.     input_names = ["input0"]
    
19.     output_names = ["output0"]
    
20.     inputs = torch.randn(1, 3, 112, 112).to(conf.device)
    
21.     # torch1.5.1 ---> onnx
    
22.     torch_out = torch.onnx._export(learner.model, inputs, output_onnx, export_params=True, verbose=False,
    
23.                                    input_names=input_names, output_names=output_names)
    
24. 
    
25. # pytorch ---> onnx
    
26. # onnx    ---> caffe2
    
27. # pip3 install onnx-caffe2
    

复制代码

import onnx
import torch
import torch.onnx


有些情况既需要torch.onnx 又需要 onnx