pytorch模型保存的格式

首先讲讲保存模型或权重参数的后缀格式,权重参数和模型参数的后缀格式一样,pytorch中最常见的模型保存使用 .pt 或者是 .pth 作为模型文件扩展名。还有其他的保存数据的格式为.t7或者.pkl格式。t7文件是沿用torch7中读取模型权重的方式,而pth文件是python中存储文件的常用格式,而在keras中则是使用.h5文件 。

我们经常会看到后缀名为.pt, .pth, .pkl的pytorch模型文件,这几种模型文件在格式上有什么区别吗?其实它们并不是在格式上有区别,只是后缀不同而已(仅此而已),在用torch.save()函数保存模型文件时,各人有不同的喜好,有些人喜欢用.pt后缀,有些人喜欢用.pth或.pkl.用相同的torch.save()语句保存出来的模型文件没有什么不同。

在pytorch官方的文档/代码里,有用.pt的,也有用.pth的。一般惯例是使用.pth,但是官方文档里貌似.pt更多,而且官方也不是很在意固定用一种。

打开pkl的方式

.pkl文件是python保存文件的一种文件格式,如果直接打开会显示一堆序列化的东西。
需要使用rb类型来打开
rb – 读取2进制文件,r – 读取文本文件

# cPickle是python2系列用的,3系列已经不用了,直接用pickle就好了
import pickle
data = pickle.load( open('a.pkl','rb'))  #记得加上'rb'
print(data)


import pickle
if __name__ == '__main__':
    a = [1, 2, 3]
    # 以二进制的形式打开文件
    with open("./1.pkl", "wb") as f:
        # 将列表a序列化后写入文件
        pickle.dump(a, f)
    with open("./1.pkl", 'rb') as f:
        # 反序列化解析成列表a
        a = pickle.load(f)
    print(a, a.__class__)  # [1, 2, 3] <class 'list'>



import os
import pickle



def read_pickle(work_path):
    data_list = []
    with open(work_path, "rb") as f:
            while True :
                try:
                    data = pickle.load(f)    
                    #data_list.append(data)
                    print(data.__class__)
                    print(data.keys())
                    print(set(data))
                    print(data['infos'].__class__)
                    print(data['metadata'].__class__)
                    print(set(data['metadata']))
                    print(data['metadata']['version'].__class__)
                    print(data['infos'][0])
                    print(data['infos'][0].__class__)
                    print(data['infos'][0].keys())
                except EOFError:
                    break
    return data_list

def main():
    pkl_path = 'nuscenes_infos_temporal_test.pkl'
    data_list = read_pickle(pkl_path)
    print(data_list)
if __name__=="__main__":
    main()

Logo

为开发者提供学习成长、分享交流、生态实践、资源工具等服务,帮助开发者快速成长。

更多推荐