簡介
Open Neural Network Exchange（ONNX，開放神經(jīng)網(wǎng)絡交換）格式，是一個用于表示深度學習模型的標準，可使模型在不同框架之間進行轉(zhuǎn)移。

ONNX是一種針對機器學習所設計的開放式的文件格式，用于存儲訓練好的模型。它使得不同的人工智能框架（如Pytorch, MXNet）可以采用相同格式存儲模型數(shù)據(jù)并交互。 ONNX的規(guī)范及代碼主要由微軟，亞馬遜 ，Facebook 和 IBM等公司共同開發(fā)，以開放源代碼的方式托管在Github上。目前官方支持加載ONNX模型并進行推理的深度學習框架有： Caffe2, PyTorch, MXNet，ML.NET，TensorRT 和 Microsoft CNTK，并且 TensorFlow也非官方的支持ONNX。
代碼實現(xiàn)( .pth --> .onnx )

import torch import torch.onnx from tinynet import tinynet from conf import settings import osdef 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 0model = 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)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)

與50位技術專家面對面20年技術見證，附贈技術全景圖

總結(jié)

以上是生活随笔為你收集整理的Pytorch模型(.pth)转onnx模型(.onnx)的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。