PyTorch学习记录(三):pth模型转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)

参考资料

posted @ 2023-03-15 10:30  达可奈特  阅读(511)  评论(0编辑  收藏  举报