from joblib.parallel import Parallel,delayed

一般用法

Joblib提供了一个简单的帮助类来编写并行化的循环。其核心思想是把代码写成生成器表达式的样子,然会再将它转换为并行计算:

from math import sqrt 
[sqrt(i ** 2) for i in range(10)]

使用以下方式,可将计算分布到两个CPU上:

from math import sqrt
from joblib import Parallel, delayed 
Parallel(n_jobs=2)(delayed(sqrt)(i ** 2) for i in range(10))

以上,Parallel对象会创建一个进程池,以便在多进程中执行每一个列表项。函数delayed是一个创建元组(function, args, kwargs)的简单技巧。

使用多线程

默认情况下,Parallel使用Python的多进程模块(multiprocessing)来fork工作进程,以便任务可以在独立的CPU上同时执行。这对于一般的Python程序来说是合理的,但这会产生一些开销,即,输入输出数据需要被序列化到一个排队,才能在工作进程之间进行通信。

当然,如果你知道,你调用的函数是基于编译扩展的,且它在执行的大部分时间都会释放Python的全局解释器锁(GIL),那么此时使用多线程可能会更高效。

为了使用多线程,只需在构造Parallel的时候设置backend='threading’即可:

Parallel(n_jobs=2, backend="threading")( ... delayed(sqrt)(i ** 2) for i in range(10))

在共享内存(memmaping)中操作数值型数据

默认情况下,当n_jobs != 1时,joblib使用Python标准库的多进程模块(multiprocessing)来创建真实的Python工作进程 。传递给Parallel调用的参数被序列化,并且会在每一个工作进程中重新创建。

这对于大型参数会成为一个问题,因为它们会被工作进程创建n_jobs次。

这在使用numpy进行科学计算中经常发生。joblib.Parallel对大型数组提供了一个特别的处理方法就是自动dump它们到文件系统,并将引用传递给工作进程,然后让工作进程使用numpy.ndarray的子类numpy.memmap以内存映射的方式打开它们 。这使得所有工作进程可以共享一段数据(更准确的说是共享一段内存)。

自动将array转换为memmap

通过在数组的大小上配置一个阀值自动触发将array转换为memmap:

import numpy as np 
from joblib import Parallel, delayed 
 from joblib.pool import has_shareable_memory 
  Parallel(n_jobs=2, max_nbytes=1e6)( ... delayed(has_shareable_memory)(np.ones(int(i))) ... for i in [1e2, 1e4, 1e6])

默认情况下,数据被dump到/dev/shm共享内存分区,如果它存在且可写 (在Linux上就是这样的)。否则将使用操作系统的临时文件夹。可以通过设置Parallel构造函数的参数temp_folder来自定义临时数据文件的位置 。

设置max_nbytes=None可禁用自动转换。

手动映射输入数据

为了更好地使用内存,你可以手动将数组dump成memmap,然后在fork工作进程之前从父进程中删除原数组。

让我们在父进程中创建一个大型数组:

large_array = np.ones(int(1e6))

然后,将它dump到本地文件,以便内存映射:

import tempfile 
import os 
from joblib import load, dump 
temp_folder = tempfile.mkdtemp() 
filename = os.path.join(temp_folder, 'joblib_test.mmap') 
if os.path.exists(filename):
      os.unlink(filename) 
_ = dump(large_array, filename) 
large_memmap = load(filename, mmap_mode='r+')

此时,变量large_memmap指向一个numpy.memmap实例:

>>> large_memmap.__class__.__name__, large_array.nbytes, large_array.shape 
>('memmap', 8000000, (1000000,)) 
>>> np.allclose(large_array, large_memmap) 
>True

然后,我们就可以释放原来的数组了:

del large_array 
import gc
_ = gc.collect()

large_memmap还可以被切片成小的memmap:

>>> small_memmap = large_memmap[2:5] 
>>> small_memmap.__class__.__name__, small_memmap.nbytes, small_memmap.shape
> ('memmap', 24, (3,))

最后,对np.ndarray视图的修改会被写回原来的内存映射文件:

>>> small_array = np.asarray(small_memmap) 
>>>> small_array.__class__.__name__, small_array.nbytes, small_array.shape 
>('ndarray', 24, (3,))

所有这三个结构都指向相同的内存区域,且这段内存能够被工作进程直接使用:

>>> Parallel(n_jobs=2, max_nbytes=None)( ... delayed(has_shareable_memory)(a) ... for a in [large_memmap, small_memmap, small_array]) 
>[True, True, True]