从源码理解pickle和joblib加载dict的性能不同

Posted 2021-06-29 ybdesire

1. 引入

最近有发现，pickle在加载（load）比较大的dict时，速度是比joblib快的。

上网查了下pickle和joblib的区别，发现写这个主题的内容比较少。

所以本文试对“pickle和joblib在加载dict时的快慢区别”这个主题进行了一些测试与研究。

2. 验证 pickle 与 joblib 加载 dict 快慢测试

使用如下代码，首先建立一个比较大的dict，并用pickle与joblib分别进行dump/load测试。

import time
import pickle
import joblib as jl

def get_big_dict():
    d = {}
    for i in range(10000000):
        key = 'k'+str(i)
        value = i
        d[key]=value


def test_pickle_dump_load_big_dict(d):
    dump_time_cost_list = []
    load_time_cost_list = []
    for i in range(100):#分别dump/load 100次，取平均值
        # dump
        t1 = time.time()
        fw = open('tmpfile.bin','wb')
        pickle.dump(d, fw)
        fw.close()
        t2 = time.time()
        # load
        fr = open('tmpfile.bin','rb')
        d2 = pickle.load(fr)
        fr.close()
        t3 = time.time()
        dump_time_cost_list.append(t2-t1)
        load_time_cost_list.append(t3-t2)
    print('pickle dump time cost ave: {0}'.format(sum(dump_time_cost_list)/len(dump_time_cost_list)))
    print('pickle load time cost ave: {0}'.format(sum(load_time_cost_list)/len(load_time_cost_list)))


def test_joblib_dump_load_big_dict(d):
    dump_time_cost_list = []
    load_time_cost_list = []
    for i in range(100):#分别dump/load 100次，取平均值
        # dump
        t1 = time.time()
        jl.dump(d, 'tmpfile.bin')
        t2 = time.time()
        # load
        d2 = jl.load('tmpfile.bin')
        t3 = time.time()
        dump_time_cost_list.append(t2-t1)
        load_time_cost_list.append(t3-t2)
    print('joblib dump time cost ave: {0}'.format(sum(dump_time_cost_list)/len(dump_time_cost_list)))
    print('joblib load time cost ave: {0}'.format(sum(load_time_cost_list)/len(load_time_cost_list)))



if __name__=='__main__':
    d = get_big_dict()
    test_pickle_dump_load_big_dict(d)
    test_joblib_dump_load_big_dict(d)

在这个代码中，首先生成一个比较大的dict。然后分别运行pickle/joblib来dump/load这个dict数据，分别进行100次的dump/load，然后计算其dump/load所用的时间，并对100次测试结果取平均值。其中一次实验得到的结果为：

pickle dump time cost ave: 0.00247844934463501
pickle load time cost ave: 0.0010942578315734862
joblib dump time cost ave: 0.006253149509429932
joblib load time cost ave: 0.0012739634513854981

在python3.6环境下，经过多次运行测试，最终结果都是在load数据时，pickle比joblib快约20%。在dump数据时，pickle也比joblib快。这个结论是能稳定重现的。

那为什么对于dict这样的数据加载，pickle会更快呢？

3. joblib加载数据的过程

下面通过源码来找到joblib加载数据的逻辑。

从参考2中，找到了joblib加载数据的入口，本文精简如下

def load(filename, mmap_mode=None):
    fobj = filename
    filename = getattr(fobj, 'name', '')
    with _read_fileobject(fobj, filename, mmap_mode) as fobj:
        obj = _unpickle(fobj)
    return obj

从这里可以看到，joblib支持memory map机制，从参考1可以发现，这使得joblib能支持多进程共享内存，pickle不具备这个能力。
加载数据的关键在于_unpickle()，进一步找到其源码（详见参考3），这里简化如下：

def _read_fileobject(fileobj, filename, mmap_mode=None):
    compressor = _detect_compressor(fileobj)

    if compressor == 'compat':
        # other logic
    else:
        if compressor in _COMPRESSORS:
            inst = compressor_wrapper.decompressor_file(fileobj)
            fileobj = _buffered_read_file(inst)
        yield fileobj

从上面的源码可以看到，在对数据进行读取前，会根据文件压缩方式，对文件进行解压。这里的关键函数是_buffered_read_file()。
其源码（详见参考4）简化如下：

def _buffered_read_file(fobj):
    """Return a buffered version of a read file object."""
    return io.BufferedReader(fobj, buffer_size=_IO_BUFFER_SIZE)

从这里可以发现，joblib最底层的数据加载，是用io这个库中的BufferedReader()函数来实现的。其用法见参考5.

4. pickle加载数据的过程

pickle加载数据的源码见参考6，简化后如下：

def load(self):
    while True:
        key = read(1)

其中read()的实现见参考7，简化后：

def read(self, size=-1):
    b = bytearray(size.__index__())
    n = self.readinto(b)
    return bytes(b)

这里没有进一步跟进readinto()，但跟进到这里，也就能发现pickle并没有像joblib一样在加载数据之前做数据解压、内存映射之类的操作。

5. 总结

本文对joblib和pickle加载dict的快慢进行了一些测试和源码分析，最终得到如下结论：

1. 在加载load数据（dict）时，pickle比joblib快约20%，每次试验基本都能重现这个结果
2. 在dump数据（dict）时，pickle比joblib快，但每次试验的数据不一样，只是快，快多少不好重现
3. 经过对加载数据的源码进行分析，发现joblib在加载数据时，会对数据做更多一些操作（如下），所以这里应该就是joblib加载数据慢的原因
   • 压缩格式判断，数据解压
   • 内存映射相关的处理

注意，本文只探索了一个主题：pickle和joblib在加载dict时的快慢区别”，本文所研究的过程与结论仅对这个主题有效。所以，对加载numpy这样的数据，或其他不同于本文的场景，可能结论会有变化。

从参考1中，也发现了如下结论

1. 对于大numpy数据的dump/load来说，joblib更快
   • 本文并未测试这个结论
2. 如果不是对大numpy数据进行dump/load，pickle可能更快
   • 本文从测试结果与源码分析都证明了这个结论

参考

1. https://stackoverflow.com/questions/12615525/what-are-the-different-use-cases-of-joblib-versus-pickle

2. https://github.com/joblib/joblib/blob/master/joblib/numpy_pickle.py#L531

3. https://github.com/joblib/joblib/blob/754433f617793bc950be40cfaa265a32aed11d7d/joblib/numpy_pickle_utils.py#L96

4. https://github.com/joblib/joblib/blob/754433f617793bc950be40cfaa265a32aed11d7d/joblib/numpy_pickle_utils.py#L85

5. https://docs.python.org/zh-cn/3/library/io.html#io.BufferedReader

6. https://github.com/python/cpython/blob/main/Lib/pickle.py#L1187

7. https://github.com/python/cpython/blob/f8a95df84bcedebc0aa7132b3d1a4e8f000914bc/Lib/_pyio.py#L646