pytorch保存提取net的不同文件格式pkl和pth

我们在训练深度学习模型的过程中,最好对已经训练好的深度学习模型进行保存,或者方便的加载别人训练好的模型微调节省训练时间,实现高效率解决问题。

为什么需要模型文件

  • 深度学习的模型参数超级多比如:Transformer模型、Bert模型等。
  • 训练的数据集一般很大,比如:1000G以上等。
  • 若本地电脑的算力或者实验室的服务器算力基本不够,训练模型花费时间多,一个模型短则训练几天不能停,甚至几个月,有可能发生内存不够等。这时若有类似的训练好的模型可以直接拿来用然后微调是非常好的。

torch.save()实现对网络结构和模型参数的保存.有两种保存方式

  • 一是保存整个神经网络的的结构信息和模型参数信息,save的对象是整个网络模型;
  • 二是只保存神经网络的训练模型参数,save的对象是net.state_dict()。

假设有一个训练好的模型名叫net1

torch.save(net1,7-net.pth’) # 保存整个神经网络的结构和模型参数

torch.save(net1,7-net.pkl’) # 同上

torch.save(net1.state_dict(),7-net_params.pth’) # 只保存神经网络的模型参数

torch.save(net1.state_dict(),7-net_params.pkl’) # 同上

如果使用torch.save方法来进行模型参数的保存,那保存文件的后缀其实没有任何影响,结果都是一样的,很多.pkl的文件也是用torch.save保存下来的,和.pth文件一模一样的

  • 不过,如果应用场景不是在这里,这两种格式的文件还是有区别的,.pkl文件是python里面保存文件的一种格式,如果直接打开会显示一堆序列化的东西,其实就是以二进制形式存储的,如果去read这些文件,则需要用’rb’而不是’r’模式。

  • 而.pth文件则有不同的应用,Python在遍历已知的库文件目录过程中,如果见到一个.pth 文件,就会将文件中所记录的路径加入到 sys.path 设置中,于是 .pth 文件指明的库也就可以被 Python 运行环境找到了。

但其实不管pkl文件还是pth文件,都是以二进制形式存储的,没有本质上的区别,你用pickle这个库去加载pkl文件或pth文件,效果都是一样的。

模型文件的后缀

  • .pt :这个后缀在官方文档里使用较多。
  • .pth :这个后缀一般大家觉得惯例使用这个。
  • .pkl:这个后缀是因为 Python 有一个序列化模块 pickle ,然后使用它保存模型时,通常会起一个以 .pkl为后缀名的文件。

猜你喜欢

转载自blog.csdn.net/weixin_45277161/article/details/128949571
今日推荐