简介
Open Neural Network Exchange(ONNX,开放神经网络交换)格式,是一个用于表示深度学习模型的标准,可使模型在不同框架之间进行转移。
ONNX是一种针对机器学习所设计的开放式的文件格式,用于存储训练好的模型。它使得不同的人工智能框架(如Pytorch, MXNet)可以采用相同格式存储模型数据并交互。 ONNX的规范及代码主要由微软,亚马逊 ,Facebook 和 IBM等公司共同开发,以开放源代码的方式托管在Github上。目前官方支持加载ONNX模型并进行推理的深度学习框架有: Caffe2, PyTorch, MXNet,ML.NET,TensorRT 和 Microsoft CNTK,并且 TensorFlow也非官方的支持ONNX。
代码实现( .pth --> .onnx )
import torch
import torch.onnx
from tinynet import tinynet
import os
def pth_to_onnx(input, checkpoint, onnx_path, input_names=['input'], output_names=['output'], device='cpu'):
if not onnx_path.endswith('.onnx'):
print('Warning! The onnx model name is not correct,\
please give a name that ends with \'.onnx\'!')
return 0
model = tinynet() #导入模型
model.load_state_dict(torch.load(checkpoint)) #初始化权重
model.eval()
# model.to(device)
torch.onnx.export(model, input, onnx_path, verbose=True, input_names=input_names, output_names=output_names) #指定模型的输入,以及onnx的输出路径
print("Exporting .pth model to onnx model has been successful!")
if __name__ == '__main__':
os.environ['CUDA_VISIBLE_DEVICES']='2'
checkpoint = './tinynet.pth'
onnx_path = './tinynet.onnx'
input = torch.randn(1, 1, 640, 360)
# device = torch.device("cuda:2" if torch.cuda.is_available() else 'cpu')
pth_to_onnx(input, checkpoint, onnx_path)