原則:
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)
參考資料
- 文檔
https://pytorch.org/docs/stable/onnx.html - 教程
https://pytorch.org/tutorials/advanced/super_resolution_with_onnxruntime.html - 代碼
https://github.com/pytorch/tutorials/blob/master/advanced_source/super_resolution_with_onnxruntime.py - PyTorch保存/加載模型
https://pytorch.org/tutorials/beginner/saving_loading_models.html - Netron Issue
Not able to visualize PyTorch model