【Python】pickle的正确使用姿势

本文已参与「新人创作礼」活动，一起开启掘金创作之路

经常使用pickle存取工作环境，时常在想，既然pickle要把数据存在文件内，为什么不把文件打开关闭的操作集成在API里呢？这篇文章就解决了这个问题。

首先简要介绍一下pickle：pickle提供了一个简单的持久化功能。可以将对象以文件的形式存放在磁盘上。pickle模块只能在python中使用，python中几乎所有的数据类型（列表，字典，集合，类等）都可以用pickle来序列化。主要用到下面两个接口：

pickle.dump(obj, file[, protocol])

序列化对象，并将结果数据流写入到文件对象中。参数protocol是序列化模式，默认值为0，表示以文本的形式序列化。protocol的值还可以是1或2，表示以二进制的形式序列化。

pickle.load(file)

反序列化对象。将文件中的数据解析为一个Python对象。

长话短说，下面两个接口，基于pickle实现了变量以文件形式存取：

import pickle as pkl

def save_file(var, filename):
    """输入一个变量和一个文件名，将这个变量存入这个文件名对应的文件中。这个函数基于pickle.dump
    @param: var(python var): python中的一个变量
    @param: filename(str): 变量将存入这个目标路径。这个路径最好是一个.pkl文件
    @returns: None
    @usage example: save_file(name_list, "save/namelist.pkl")
    """
    with open(filename, 'wb') as f:
        pkl.dump(var, f)


def load_file(filename):
    """读取一个文件名，加载这个文件名对应的内容。这个函数基于pickle.load
    @param: filename(str): 文件名，对应一个通过save_file生成的文件
    @returns: var(python var: 这个文件对应的python变量
    @usage example: name_list = open_file("save/namelist.pkl")
    """
    with open(filename, 'rb') as f:
        var = pkl.load(f)
    return var

提供对应的单元测试（或者称为使用范例）供君参考：

 def test_save_file():
    va = [2, 3, 5, 7, 11]
    save_file(va, "ls.pkl")


def test_load_file():
    va = load_file('ls.pkl')
    print(va)

这两个接口适用于一切类型的变量。灵活使用这两个接口，可以在python中实现类似matlab workspace的功能。

常见错误

不建议以此法存取一些依赖特定结构的东西。以pytorch的模型为例，如果以此法存取pytorch，要求存取之处代码环境的import环境完全一样。类似地，MATLAB的workspace也有相应的问题。