古月居 PyTorch入门:一起从零搭建神经网络 七、PyTorch车牌字符识别项目四
1、如何把训练好的模型打包成文件
我们在搭建好数据集和神经网络之后,每次运行如果都提前训练,那将会花费相当长的时间。所以我们就需要将训练好的模型参数进行保存。这样做的好处是不用每次都训练,可以直接调用训练好的文件包。
对于模型的保存,一共有两种方式,一个是保存整个模型,另一种是只保存模型的参数(常用)。
函数解释:
代码:
torch.save(cnn,'CNN.pth')#保存模型
torch.save(cnn.state_dict(),'CNN.pth')#保存模型参数
cnn =CNN_h.CNN()
cnn.load_state_dict(torch.load('CNN.pth'))#加载模型
代码全集:添加链接描述
2、结果测试
运行正确率大概在50%-80%。