东方耀AI技术分享
标题:
pytorch模型转换为onnx
[打印本页]
作者:
东方耀
时间:
2020-8-3 17:50
标题:
pytorch模型转换为onnx
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
欢迎光临 东方耀AI技术分享 (http://www.ai111.vip/)
Powered by Discuz! X3.4