PyTorch學習記錄(三):onnx模型部署

原則:

Python for Training
C++ for Inference

PyTorch模型導出:
PyTorch使用.pth文件來對神經網絡的權重進行保存,.pth文件中的模型權重則是按照字典格式進行保存的,但是.pth文件中沒有網絡的結果信息。需要藉助開放神經網絡交換(Open Neural Network Exchange, ONNX)框架將模型導出爲結構和權重完整的.onnx文件。

模型參數量

model = FPN()
num_params = sum(p.numel() for p in model.parameters())
print("num of params: {:.2f}k".format(num_params/1000.0))
# torch.numel()返回tensor的元素數目,即number of elements

打印模型

model = FPN()
num_params = sum(p.numel() for p in model.parameters())
print("num of params: {:.2f}k".format(num_params/1000.0))
print("===========================")
#for p in model.parameters():
#    print(p.name)
print(model)

參考資料

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章