|
pytorch模型转换为onnx
- import torch.onnx
- from config import get_config_lite
- from Learner import face_learner
- if __name__ == '__main__':
- # input_model = "./work_space_03_lite/save/model_2020-07-18-02-41_accuracy:0.9261636148230064_step:3311920_final.pth"
- output_onnx = './work_space_03_lite/save/FaceRecog_dfy.onnx'
- print("==> Exporting model to ONNX format at '{}'".format(output_onnx))
- conf = get_config_lite(training=False)
- # conf = get_config_dfy(training=False)
- print("easydict的配置参数:", conf)
- learner = face_learner(conf, inference=True)
- fixed_str = "2020-07-18-02-41_accuracy:0.9261636148230064_step:3311920_final.pth"
- learner.load_state(conf, fixed_str, model_only=True, from_save_folder=True)
- # 开始:pytorch模型转换为onnx
- input_names = ["input0"]
- output_names = ["output0"]
- inputs = torch.randn(1, 3, 112, 112).to(conf.device)
- # torch1.5.1 ---> onnx
- torch_out = torch.onnx._export(learner.model, inputs, output_onnx, export_params=True, verbose=False,
- input_names=input_names, output_names=output_names)
- # pytorch ---> onnx
- # onnx ---> caffe2
- # pip3 install onnx-caffe2
复制代码
import onnx
import torch
import torch.onnx
有些情况既需要torch.onnx 又需要 onnx
|
|