文章目录
1. .pth、.pt、.pkl以及.state等文件有什么区别?
通常,在神经网络的训练过程中,会产生.pth、.pt、.pkl以及.state等文件。使用不同的后缀对于保存什么样的模型是没有影响的。在torch.save()函数中,torch只是调用了Python的pickle来完成,而且没有做任何的后缀名判断,因此无论保存成什么后缀都可以。其中保存的都是模型+模型权重参数或者单独的模型权重参数。
2. 两种不同的保存方式
2.1 通过state_dict参数来保存和读取模型的权重参数
其后缀推荐使用.pt或者.pth,但是没有硬性规定。
需要注意的是,load_state_dict()函数需要传进的参数是一个字典对象,并非是PATH对象,因此需要用torch.load(PATH)对其进行反序列化。即不能直接使用load_state_dict(PATH)。对比上面给出的代码。