[python] Python平行計算庫Joblib使用指北

Joblib是用於高效平行計算的Python開源庫，其提供了簡單易用的記憶體對映和平行計算的工具，以將任務分發到多個工作程序中。Joblib庫特別適合用於需要進行重複計算或大規模資料處理的任務。Joblib庫的官方倉庫見：joblib，官方文件見：joblib-doc。

Jolib庫安裝程式碼如下：

pip install joblib

# 檢視版本
import joblib
joblib.__version__

'1.4.2'

目錄

• 1 使用說明
  • 1.1 Memory類
  • 1.2 Parallel類
  • 1.3 序列化
• 2 例項
  • 2.1 joblib快取和並行
  • 2.2 序列化
  • 2.3 記憶體監視
• 3 參考

1 使用說明

Joblib庫主要功能涵蓋以下三大塊：

• 記憶模式：Memory類將函式的返回值快取到磁碟。下次呼叫時，如果輸入引數不變，就直接從快取中載入結果，避免重複計算。
• 平行計算：Parallel類將任務拆分到多個程序或者執行緒中並行執行，加速計算過程。
• 高效的序列化：針對NumPy陣列等大型資料物件進行了最佳化，且序列化和反序列化速度快。

1.1 Memory類

Joblib庫的Memory類支援透過記憶模式，將函式的計算結果儲存起來，以便在下次使用時直接呼叫。這種機制的優勢在於加速計算過程、節約資源以及簡化管理。

Memory類建構函式如下：

class joblib.Memory(location=None, backend='local', mmap_mode=None, compress=False, verbose=1, bytes_limit=None, backend_options=None)

引數介紹如下：

• location: 快取檔案的存放位置。如果設定為 None，則不快取。
• backend: 快取的後端儲存方式。預設是 "local"，表示使用本地檔案系統。
• mmap_mode: 一個字串，表示記憶體對映檔案的模式（None, ‘r+’, ‘r’, ‘w+’, ‘c’）。
• compress: 表示是否壓縮快取檔案。壓縮可以節省磁碟空間，但會增加 I/O 操作的時間。
• verbose: 一個整數，表示日誌的詳細程度。0 表示沒有輸出，1 表示只輸出警告，2 表示輸出資訊，3 表示輸出除錯資訊。
• bytes_limit: 一個整數或 None，表示快取使用的位元組數限制。如果快取超過了這個限制，最舊的快取檔案將被刪除。
• backend_options: 傳遞給快取後端的選項。

Memory類簡單使用

下面程式碼展示第一次呼叫函式並快取結果：

from joblib import Memory
import os, shutil
# 建立一個Memory物件，指定快取目錄為當前目錄下的run資料夾
# verbose=0表示關閉詳細輸出
cachedir = './run'
if os.path.exists(cachedir):
    shutil.rmtree(cachedir)
memory = Memory(cachedir, verbose=0)

# 使用@memory.cache裝飾器，將函式f的結果快取起來
@memory.cache
def f(x):
    # 只有當函式的輸入引數x沒有被快取時，才會執行函式體內的程式碼
    print('Running f(%s)' % x)
    return x

# 第一次呼叫f(1)，會執行函式體內的程式碼，並將結果快取起來
print(f(1))

Running f(1)
1

第二次呼叫函式：

# 第二次呼叫f(1)，由於結果已經被快取，不會再次執行函式體內的程式碼，而是直接從快取中讀取結果
print(f(1))

1

呼叫其他函式：

# 呼叫f(2)，由於輸入引數不同，會再次執行函式體內的程式碼，並將結果快取起來
print(f(2))

Running f(2)
2

將Memory類應用於numpy陣列

import numpy as np
from joblib import Memory
import os, shutil
cachedir = './run'
if os.path.exists(cachedir):
    shutil.rmtree(cachedir)
memory = Memory(cachedir, verbose=0)

@memory.cache
def g(x):
    print('A long-running calculation, with parameter %s' % x)
    # 返回漢明窗
    return np.hamming(x)

@memory.cache
def h(x):
    print('A second long-running calculation, using g(x)')
    # 生成範德蒙德矩陣
    return np.vander(x)

# 呼叫函式g，傳入引數3，並將結果儲存在變數a中
a = g(3)
# 列印變數a的值
print(a)

# 再次呼叫函式g，傳入相同的引數3，由於結果已被快取，不會重新計算
print(g(3))

A long-running calculation, with parameter 3
[0.08 1.   0.08]
[0.08 1.   0.08]

直接計算和快取結果是等同的：

# 呼叫函式h，傳入變數a作為引數，並將結果儲存在變數b中
b = h(a)
# 再次呼叫函式h，傳入相同的引數a，由於結果已被快取，不會重新計算
b2 = h(a)

# 使用numpy的allclose函式檢查b和b2是否足夠接近，即它們是否相等
print(np.allclose(b, b2))

A second long-running calculation, using g(x)
True

直接呼叫快取結果

import numpy as np
from joblib import Memory

import os, shutil

# 設定快取目錄的路徑。
cachedir = './run'

# 檢查快取目錄是否存在。
if os.path.exists(cachedir):
    # 如果快取目錄存在，使用shutil.rmtree刪除該目錄及其內容。
    shutil.rmtree(cachedir)

# 初始化Memory物件，設定快取目錄為上面定義的cachedir，mmap_mode設定為'r'，表示只讀模式。
memory = Memory(cachedir, mmap_mode='r', verbose=0)

# 使用memory.cache裝飾器快取np.square函式的結果。
square = memory.cache(np.square)

a = np.vander(np.arange(3)).astype(float)

# 列印透過square函式處理後的矩陣a。
print(square(a))

# 獲取a的快取結果
result = square.call_and_shelve(a)
print(result.get())  # 獲取並列印快取的結果。

[[ 0.  0.  1.]
 [ 1.  1.  1.]
 [16.  4.  1.]]
[[ 0.  0.  1.]
 [ 1.  1.  1.]
 [16.  4.  1.]]

類中使用快取

Memory類不建議將其直接用於類方法。如果想在類中使用快取，建議的模式是在類中使用單獨定義的快取函式，如下所示：

@memory.cache
def compute_func(arg1, arg2, arg3):
    pass

class Foo(object):
    def __init__(self, args):
        self.data = None

    def compute(self):
        # 類中呼叫快取的函式
        self.data = compute_func(self.arg1, self.arg2, 40)

1.2 Parallel類

Joblib庫的Parallel類用於簡單快速將任務分解為多個子任務，並分配到不同的CPU核心或機器上執行，從而顯著提高程式的執行效率。

Parallel類建構函式及主要引數如下：

class joblib.Parallel(n_jobs=default(None), backend=default(None), return_as='list', verbose=default(0), timeout=None, batch_size='auto', pre_dispatch='2 * n_jobs', temp_folder=default(None), max_nbytes=default('1M'), require=default(None))

引數介紹如下：

• n_jobs: 指定並行任務的數量，為-1時表示使用所有可用的CPU核心；為None時表示使用單個程序。
• backend：指定並行化的後端，可選項：
  • 'loky'：使用loky庫實現多程序，該庫由joblib開發者開發，預設選項。
  • 'threading'：使用threading庫實現多執行緒。
  • 'multiprocessing'：使用multiprocessing庫實現多程序。
• return_as：返回結果格式，可選項：
  • 'list：列表。
  • generator：按照任務提交順序生成結果的生成器。
  • generator_unordered：按照執行結果完成先後順序的生成器。
• verbose: 一個整數，表示日誌的詳細程度。0 表示沒有輸出，1 表示只輸出警告，2 表示輸出資訊，3 表示輸出除錯資訊。
• timeout：單個任務最大執行時長，超時將引發TimeOutError。僅適用於n_jobs不為1的情況。
• batch_size：當Parallel類執行任務時，會將任務分批處理。batch_size引數決定了每個批次中包含的任務數。
• pre_dispatch: 用來決定在平行計算開始之前，每個批次有多少個任務會被預先準備好並等待被分配給單個工作程序。預設值為“2*n_jobs”，表示平行計算時可以使用2倍工作程序的任務數量。
• temp_folder：指定臨時檔案的儲存路徑。
• max_nbytes：傳遞給工作程式的陣列大小的閾值。
• require：對執行任務的要求，可選None和sharedmem。sharedmem表示將使用共享記憶體來執行並行任務，但會影響計算效能。

簡單示例

以下程式碼展示了單執行緒直接執行計算密集型任務結果：

from joblib import Parallel, delayed
import numpy as np
import time

start = time.time()

# 定義一個計算密集型函式
def compute_heavy_task(data):
    # 模擬處理時間
    time.sleep(1)
    # 數值計算
    result = np.sum(np.square(data))
    return result

# 生成一些模擬資料
# 設定隨機數生成器的種子
np.random.seed(42) 
data = np.random.rand(10, 1000)  # 10個1000維的向量
results = [compute_heavy_task(d) for d in data]

# 列印結果的和
print(f"結果: {sum(results)}")
print(f"耗時：{time.time()-start}s")

結果: 3269.16485027708
耗時：10.101513624191284s

以下程式碼展示利用Parallel類建立多程序執行計算密集型任務結果：

from joblib import Parallel, delayed
import numpy as np
import time

start = time.time()

# 定義一個計算密集型函式
def compute_heavy_task(data):
    # 模擬處理時間
    time.sleep(1)
    # 數值計算
    result = np.sum(np.square(data))
    return result
    
# 設定隨機數生成器的種子
np.random.seed(42) 
# 生成一些模擬資料
data = np.random.rand(10, 1000)  # 10個1000維的向量

# 使用Parallel來並行執行任務
results = Parallel(n_jobs=8, return_as="generator")(delayed(compute_heavy_task)(d) for d in data)

# 列印結果的和
print(f"結果: {sum(results)}")
print(f"耗時：{time.time()-start}s")

結果: 3269.16485027708
耗時：2.381772041320801s

可以看到joblib庫利用多程序技術顯著提高了任務執行的效率。然而，當面對I/O密集型任務或執行時間極短的任務時，多執行緒或多程序的優勢可能並不明顯。這是因為執行緒建立和上下文切換的開銷有時可能超過任務本身的執行時間。以上述的compute_heavy_task函式為例，如果移除了其中的time.sleep函式，多程序執行所需的時間將會顯著增加。

此外獲取當前系統的cpu核心數（邏輯處理器）程式碼如下：

import joblib

# 獲取當前系統的cpu核心數
n_cores = joblib.cpu_count()

print(f'系統的核心數是：{n_cores}')

系統的核心數是：16

不同並行方式對比

以下程式碼展示了不同並行方式在Parallel類中的應用。預設使用loky多程序：

# 使用loky多程序
from joblib import Parallel, delayed
import numpy as np
import time

start = time.time()

# 定義一個計算密集型函式
def compute_heavy_task(data):
    # 模擬處理時間
    time.sleep(1)
    # 數值計算
    result = np.sum(np.square(data))
    return result

# 生成一些模擬資料
data = np.random.rand(10, 1000)  # 10個1000維的向量
results = Parallel(n_jobs=8, return_as="generator", backend='loky')(delayed(compute_heavy_task)(d) for d in data)

# 列印結果的和
print(f"結果: {sum(results)}")
print(f"耗時：{time.time()-start}s")

結果: 3382.3336437893217
耗時：2.042675256729126s

以下程式碼展示了threading多執行緒的使用，注意由於Python的全域性直譯器鎖（GIL）確保在任何時刻只有一個執行緒執行Python位元組碼。這表明即使在多核處理器上，Python的執行緒也無法實現真正的平行計算。然而，當涉及到處理I/O密集型任務或需要快速響應的小規模任務時，多執行緒依然具有優勢：

# 使用threading多執行緒
start = time.time()
results = Parallel(n_jobs=8, return_as="generator", backend = 'threading')(delayed(compute_heavy_task)(d) for d in data)

# 列印結果的和
print(f"結果: {sum(results)}")
print(f"耗時：{time.time()-start}s")

結果: 3382.3336437893217
耗時：2.040527105331421s

以下程式碼展示了multiprocessing多程序的使用，注意Windows下需要將multiprocessing相關程式碼放在main函式中：

from joblib import Parallel, delayed
import numpy as np
import time

# 定義一個計算密集型函式
def compute_heavy_task(data):
    # 模擬處理時間
    time.sleep(1)
    # 數值計算
    result = np.sum(np.square(data))
    return result

def main():
    start = time.time()

    # 生成一些模擬資料
    data = np.random.rand(10, 1000)  # 10個1000維的向量
    # multiprocessing不支援返回rgenerator
    results = Parallel(n_jobs=8, return_as="list",  backend='multiprocessing')(delayed(compute_heavy_task)(d) for d in data)

    # 列印結果的和
    print(f"結果: {sum(results)}")
    print(f"耗時：{time.time()-start}s")
    
if __name__ == '__main__':
    main()

結果: 3304.6651996375645
耗時：2.4303956031799316s

以下是loky、threading和 multiprocessing的一些關鍵特性對比：

特性/庫	loky	threading	multiprocessing
適用平臺	跨平臺	跨平臺	跨平臺，但Windows上存在限制
程序/執行緒模型	程序	執行緒	程序
GIL影響	無	有	無
適用場景	CPU密集型任務	I/O密集型任務	CPU密集型任務
啟動開銷	較小	較小	較大
記憶體使用	較高	較低	較高
程序間通訊	透過管道、佇列等	透過共享資料結構	透過管道、佇列等
執行緒間通訊	共享資料結構	共享資料結構	不適用
異常處理	程序間獨立	執行緒間共享	程序間獨立
除錯難度	較高	較低	較高
適用框架	通用	通用	通用

Python中執行緒和程序簡單對比如下：

• 資源共享：執行緒共享同一程序的記憶體和資源，而程序擁有獨立的記憶體空間。
• GIL影響：執行緒受GIL限制，程序不受GIL限制。
• 開銷：執行緒的建立和切換開銷小，程序的建立和切換開銷大。
• 適用性：執行緒適合I/O密集型任務，程序適合CPU密集型任務。
• 通訊：執行緒間通訊簡單但需要處理同步問題，程序間通訊複雜但天然隔離。

在實際應用中，選擇使用執行緒還是程序取決於任務的特性和效能需求。如果任務主要是I/O密集型，使用執行緒可以提高效能；如果任務是CPU密集型，使用程序可以更好地利用多核處理器的計算能力。

共享記憶體

預設情況下，Parallel類執行任務時各個任務不共享記憶體，如下所示：

from joblib import Parallel, delayed
shared_set = set()
def collect(x):
   shared_set.add(x)
Parallel(n_jobs=2)(delayed(collect)(i) for i in range(5))
print(sorted(shared_set))

[]

透過設定require='sharedmem'可以實現記憶體共享：

# require='sharedmem'表示需要共享記憶體，以確保多個程序可以訪問shared_set集合
Parallel(n_jobs=2, require='sharedmem')(delayed(collect)(i) for i in range(5))
print(sorted(shared_set))

[0, 1, 2, 3, 4]

上下文管理器

一些演算法需要對一個並行函式進行多次連續呼叫，但在迴圈中多次呼叫joblib.Parallel是次優的，因為這將多次建立和銷燬一組工作程序，從而導致顯著的效能開銷。

對於這種情況，使用joblib.Parallel類的上下文管理器API更為高效，可以重用同一組工作程序進行多次呼叫joblib.Parallel物件。如下所示：

from joblib import Parallel, delayed
import math
with Parallel(n_jobs=2) as parallel:
   accumulator = 0.
   n_iter = 0
   while accumulator < 1000:
       results = parallel(delayed(math.sqrt)(accumulator + i ** 2) for i in range(5))
       accumulator += sum(results)
       n_iter += 1
print(accumulator, n_iter)  

1136.5969161564717 14

parallel_config

Joblib提供parallel_config類用於配置並行執行的引數，比如並行的後端型別、批處理大小等，這些配置可以影響後續所有的parallel例項。它通常在呼叫Parallel類之前使用。關於parallel_config使用見：parallel_config。

1.3 序列化

joblib.dump()和joblib.load()提供了一種替代pickle庫的方法，可以高效地序列化處理包含大量資料的任意Python物件，特別是大型的NumPy陣列。關於pickle庫使用見：Python資料序列化模組pickle使用筆記 。兩者效果對比見：

特點	pickle	joblib
效能	一般	針對NumPy陣列等大資料型別有最佳化，通常更快
並行處理	不支援	內建並行處理功能，可以加速任務
記憶體對映	不支援	支援記憶體對映，可以高效處理大檔案
壓縮	支援	支援壓縮，可以減少儲存空間
附加功能	少	提供了一些額外的功能，如快取、延遲載入等

以下程式碼展示了joblib.dump的基本使用：

from tempfile import mkdtemp

# 使用mkdtemp建立一個臨時目錄，並將目錄路徑儲存在變數savedir中。
savedir = mkdtemp(dir='./')

import os
# 檔案儲存路徑
filename = os.path.join(savedir, 'test.joblib')

import numpy as np
import pandas as pd
import joblib

# 建立一個要持久化的字典
to_persist = [('a', [1, 2, 3]), ('b', np.arange(10)), ('c', pd.DataFrame(np.ones((5,5))))]

# 使用joblib.dump函式將to_persist字典序列化並儲存到filename指定的檔案中
# 注意pickle庫無法序列化numpy資料
joblib.dump(to_persist, filename)

['./tmp82ms1z5w\\test.joblib']

使用joblib.load函式從指定的檔案中載入之前儲存的序列化資料：

joblib.load(filename)

[('a', [1, 2, 3]),
 ('b', array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])),
 ('c',
       0    1    2    3    4
  0  1.0  1.0  1.0  1.0  1.0
  1  1.0  1.0  1.0  1.0  1.0
  2  1.0  1.0  1.0  1.0  1.0
  3  1.0  1.0  1.0  1.0  1.0
  4  1.0  1.0  1.0  1.0  1.0)]

joblib.dump和joblib.load函式還接受檔案物件：

with open(filename, 'wb') as fo:  
    # 使用joblib將物件to_persist序列化並寫入檔案
   joblib.dump(to_persist, fo)
with open(filename, 'rb') as fo:  
   joblib.load(fo)

此外joblib.dump也支援設定compress引數以實現資料壓縮：

# compress引數為壓縮級別，取值為0到9，值越大壓縮效果越好。為0時表示不壓縮，預設值為0
joblib.dump(to_persist, filename, compress=1)

['./tmp82ms1z5w\\test.joblib']

預設情況下，joblib.dump使用zlib壓縮方法，因為它在速度和磁碟空間之間實現了最佳平衡。其他支援的壓縮方法包括“gzip”、“bz2”、“lzma”和“xz”。compress引數輸入帶有壓縮方法和壓縮級別就可以選擇不同壓縮方法：

joblib.dump(to_persist, filename + '.gz', compress=('gzip', 3))  
joblib.load(filename + '.gz')
joblib.dump(to_persist, filename + '.bz2', compress=('bz2', 5))  
joblib.load(filename + '.bz2')

[('a', [1, 2, 3]),
 ('b', array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])),
 ('c',
       0    1    2    3    4
  0  1.0  1.0  1.0  1.0  1.0
  1  1.0  1.0  1.0  1.0  1.0
  2  1.0  1.0  1.0  1.0  1.0
  3  1.0  1.0  1.0  1.0  1.0
  4  1.0  1.0  1.0  1.0  1.0)]

除了預設壓縮方法，lz4壓縮演算法也可以用於資料壓縮。前提是需要安裝lz4壓縮庫：

pip install lz4

在這些壓縮方法中，lz4和預設方法效果較好。lz4使用方式與其他壓縮方式一樣：

joblib.dump(to_persist, filename, compress=('lz4', 3))  

['./tmp82ms1z5w\\test.joblib']

2 例項

2.1 joblib快取和並行

本例項展示了利用joblib快取和並行來加速任務執行。以下程式碼展示了一個高耗時任務：

# 匯入time模組，用於實現延時功能
import time

# 定義一個模擬耗時計算的函式
def costly_compute(data, column):
    # 休眠1秒，模擬耗時操作
    time.sleep(1)
    # 返回傳入資料的指定列
    return data[column]

# 定義一個計算資料列平均值的函式
def data_processing_mean(data, column):
    # 呼叫costly_compute函式獲取指定列的資料
    column_data = costly_compute(data, column)
    # 計算並返回該列資料的平均值
    return column_data.mean()

# 匯入numpy庫，並設定隨機數生成器的種子，以保證結果的可復現性
import numpy as np
rng = np.random.RandomState(42)
# 生成1000行4列的隨機資料矩陣
data = rng.randn(int(1000), 4)

# 記錄開始時間
start = time.time()
# 對資料的每一列計算平均值，並將結果儲存在results列表中
results = [data_processing_mean(data, col) for col in range(data.shape[1])]
# 記錄結束時間
stop = time.time()

# 列印處理過程的描述資訊
print('\nSequential processing')
# 列印整個處理過程的耗時
print('Elapsed time for the entire processing: {:.2f} s'.format(stop - start))

Sequential processing
Elapsed time for the entire processing: 4.05 s

下段程式碼演示瞭如何使用joblib庫來快取和並行化計算上述任務：

# 匯入time模組，用於模擬耗時操作。
import time

# 定義一個使用快取的函式，用於計算資料的均值。
def data_processing_mean_using_cache(data, column):
    return costly_compute_cached(data, column).mean()

# 從joblib庫匯入Memory類，用於快取函式的輸出。
from joblib import Memory

# 設定快取的儲存位置和詳細程度
location = './cachedir'
memory = Memory(location, verbose=0)

# 使用Memory物件的cache方法來快取costly_compute函式的輸出。
costly_compute_cached = memory.cache(costly_compute)

# 從joblib庫匯入Parallel和delayed類，用於並行執行函式。
from joblib import Parallel, delayed

# 記錄開始時間。
start = time.time()

# 使用Parallel類並行執行data_processing_mean_using_cache函式，對資料的每一列進行處理。
results = Parallel(n_jobs=2)(
    delayed(data_processing_mean_using_cache)(data, col)  
    for col in range(data.shape[1])) 

# 記錄結束時間。
stop = time.time()

# 列印第一輪處理的耗時資訊，包括快取資料的時間。
print('\nFirst round - caching the data')
print('Elapsed time for the entire processing: {:.2f} s'.format(stop - start))

First round - caching the data
Elapsed time for the entire processing: 2.05 s

再次執行相同的過程，可以看到結果被快取而不是重新執行函式：

start = time.time()
results = Parallel(n_jobs=2)(
    delayed(data_processing_mean_using_cache)(data, col)
    for col in range(data.shape[1]))
stop = time.time()

print('\nSecond round - reloading from the cache')
print('Elapsed time for the entire processing: {:.2f} s'.format(stop - start))

# 如果不想使用快取結果，可以清除快取資訊
memory.clear(warn=False)

Second round - reloading from the cache
Elapsed time for the entire processing: 0.02 s

2.2 序列化

以下示例展示了在joblib.Parallel中使用序列化記憶體對映（numpy.memmap）。記憶體對映可以將大型資料集分割成小塊，並在需要時將其載入到記憶體中。這種方法可以減少記憶體使用，並提高處理速度。

定義耗時函式：

import numpy as np

data = np.random.random((int(1e7),))
window_size = int(5e5)
slices = [slice(start, start + window_size)
          for start in range(0, data.size - window_size, int(1e5))]

import time

def slow_mean(data, sl):
    time.sleep(0.01)
    return data[sl].mean()

以下程式碼是直接呼叫函式的執行結果：

tic = time.time()
results = [slow_mean(data, sl) for sl in slices]
toc = time.time()
print('\nElapsed time computing the average of couple of slices {:.2f} s'.format(toc - tic))

Elapsed time computing the average of couple of slices 1.49 s

以下程式碼是呼叫Parallel類2個程序執行的結果，由於整體任務計算耗時較少。所以Parallel類平行計算並沒有比直接呼叫函式有太多速度優勢，因為程序啟動銷燬需要額外時間：

from joblib import Parallel, delayed

tic = time.time()
results = Parallel(n_jobs=2)(delayed(slow_mean)(data, sl) for sl in slices)
toc = time.time()
print('\nElapsed time computing the average of couple of slices {:.2f} s'.format(toc - tic))

Elapsed time computing the average of couple of slices 1.00 s

以下程式碼提供了joblib.dump和load函式加速資料讀取。其中dump函式用於將data物件序列化並儲存到磁碟上的檔案中，同時建立了一個記憶體對映，使得該檔案可以像記憶體陣列一樣被訪問。當程式再次載入這個檔案時，可以使用load函式以記憶體對映模式開啟：

import os 
from joblib import dump, load  # 從joblib庫匯入dump和load函式，用於建立和載入記憶體對映檔案

# 設定記憶體對映檔案的資料夾路徑
folder = './memmap'
os.makedirs(folder, exist_ok = True)

# 將記憶體對映檔案的名稱與路徑結合
data_filename_memmap = os.path.join(folder, 'data_memmap.joblib')

# 使用dump函式將資料物件'data'儲存到記憶體對映檔案
dump(data, data_filename_memmap)

# 使用load函式載入記憶體對映檔案，mmap_mode='r'表示以只讀模式開啟
data_ = load(data_filename_memmap, mmap_mode='r')

# 記錄開始時間
tic = time.time()
results = Parallel(n_jobs=2)(delayed(slow_mean)(data_, sl) for sl in slices)

# 記錄結束時間
toc = time.time()
print('\nElapsed time computing the average of couple of slices {:.2f} s\n'.format(toc - tic))  

import shutil
# 結束時刪除對映檔案
try:
    shutil.rmtree(folder)
except: 
    pass

Elapsed time computing the average of couple of slices 0.77 s

2.3 記憶體監視

本例項展示不同並行方式的記憶體消耗情況。

建立記憶體監視器

from psutil import Process
from threading import Thread

class MemoryMonitor(Thread):
    """在單獨的執行緒中監控記憶體使用情況（以MB為單位）。"""
    def __init__(self):
        super().__init__()  # 呼叫父類Thread的建構函式
        self.stop = False  # 用於控制執行緒停止的標記
        self.memory_buffer = []  # 用於儲存記憶體使用記錄的列表
        self.start()  # 啟動執行緒

    def get_memory(self):
        """獲取程序及其子程序的記憶體使用情況。"""
        p = Process()  # 獲取當前程序
        memory = p.memory_info().rss  # 獲取當前程序的記憶體使用量
        for c in p.children():  # 遍歷所有子程序
            memory += c.memory_info().rss  # 累加子程序的記憶體使用量
        return memory

    def run(self):
        """執行緒執行的主體方法，週期性地記錄記憶體使用情況。"""
        memory_start = self.get_memory()  # 獲取初始記憶體使用量
        while not self.stop:  # 當未設定停止標記時迴圈
            self.memory_buffer.append(self.get_memory() - memory_start)  # 記錄當前記憶體使用量與初始記憶體使用量的差值
            time.sleep(0.2)  # 休眠0.2秒

    def join(self):
        """重寫join方法，設定停止標記並等待執行緒結束。"""
        self.stop = True  # 設定停止標記
        super().join()  # 呼叫父類方法等待執行緒結束

並行任務

結果返回list的並行任務：

import time
import numpy as np

def return_big_object(i):
    """生成並返回一個大型NumPy陣列物件。"""
    time.sleep(.1)  # 休眠0.1秒模擬耗時操作
    return i * np.ones((10000, 200), dtype=np.float64) 

def accumulator_sum(generator):
    """累加生成器生成的所有值，並列印進度。"""
    result = 0
    for value in generator:
        result += value
        # print(".", end="", flush=True)  # 列印點號並重新整理輸出
    # print("")  # 列印換行符
    return result

from joblib import Parallel, delayed

monitor = MemoryMonitor()  # 建立記憶體監控器例項，並啟動監視
print("Running tasks with return_as='list'...")  # 列印啟動任務資訊
res = Parallel(n_jobs=2, return_as="list")(
    delayed(return_big_object)(i) for i in range(150)  # 使用joblib的Parallel功能並行執行任務
)
res = accumulator_sum(res)  # 累加結果
print('All tasks completed and reduced successfully.')  # 列印任務完成資訊

# 報告記憶體使用情況
del res  # 清理結果以避免記憶體邊界效應
monitor.join()  # 等待記憶體監控執行緒結束
peak = max(monitor.memory_buffer) / 1e9  # 計算峰值記憶體使用量，並轉換為GB
print(f"Peak memory usage: {peak:.2f}GB")  # 列印峰值記憶體使用量

Running tasks with return_as='list'...
All tasks completed and reduced successfully.
Peak memory usage: 2.44GB

如果改為輸出生成器，那麼記憶體使用量將會大大減少：

monitor_gen = MemoryMonitor()  # 建立記憶體監控器例項，並啟動監視
print("Running tasks with return_as='generator'...")  # 列印啟動任務資訊
res = Parallel(n_jobs=2, return_as="generator")(
    delayed(return_big_object)(i) for i in range(150)  
)
res = accumulator_sum(res)  # 累加結果
print('All tasks completed and reduced successfully.')  # 列印任務完成資訊

# 報告記憶體使用情況
del res  # 清理結果以避免記憶體邊界效應
monitor_gen.join()  # 等待記憶體監控執行緒結束
peak = max(monitor_gen.memory_buffer) / 1e9  # 計算峰值記憶體使用量，並轉換為GB
print(f"Peak memory usage: {peak:.2f}GB")  # 列印峰值記憶體使用量

Running tasks with return_as='generator'...
All tasks completed and reduced successfully.
Peak memory usage: 0.19GB

下圖展示了以上兩種方法的記憶體消耗情況，第一種情況涉及到將所有結果儲存在記憶體中，直到處理完成，這可能導致記憶體使用量隨著時間線性增長。而第二種情況generator則涉及到流式處理，即結果被實時處理，因此不需要同時在記憶體中儲存所有結果，從而減少了記憶體使用的需求：

import matplotlib.pyplot as plt
plt.figure(0)
plt.semilogy(
    np.maximum.accumulate(monitor.memory_buffer),
    label='return_as="list"'
)
plt.semilogy(
    np.maximum.accumulate(monitor_gen.memory_buffer),
    label='return_as="generator"'
)
plt.xlabel("Time")
plt.xticks([], [])
plt.ylabel("Memory usage")
plt.yticks([1e7, 1e8, 1e9], ['10MB', '100MB', '1GB'])
plt.legend()
plt.show()

進一步節省記憶體

前一個例子中的生成器是保持任務提交的順序的。如果某些程序任務提交晚，但比其他任務更早完成。相應的結果會保持在記憶體中，以等待其他任務完成。如果任務對結果返回順序無要求，例如最後只是對所有結果求和，可以使用generator_unordered減少記憶體消耗。如下所示：

# 建立一個每個任務耗時可能不同的處理函式
def return_big_object_delayed(i):
    if (i + 20) % 60:
        time.sleep(0.1)
    else:
        time.sleep(5)
    return i * np.ones((10000, 200), dtype=np.float64)

返回為generator格式的記憶體使用：

monitor_delayed_gen = MemoryMonitor()
print("Create result generator on delayed tasks with return_as='generator'...")
res = Parallel(n_jobs=2, return_as="generator")(
    delayed(return_big_object_delayed)(i) for i in range(150)
)
res = accumulator_sum(res)
print('All tasks completed and reduced successfully.')

del res 
monitor_delayed_gen.join()
peak = max(monitor_delayed_gen.memory_buffer) / 1e6
print(f"Peak memory usage: {peak:.2f}MB")

Create result generator on delayed tasks with return_as='generator'...
All tasks completed and reduced successfully.
Peak memory usage: 784.23MB

返回為generator_unordered格式的記憶體使用：

monitor_delayed_gen_unordered = MemoryMonitor()
print(
  "Create result generator on delayed tasks with "
  "return_as='generator_unordered'..."
)
res = Parallel(n_jobs=2, return_as="generator_unordered")(
    delayed(return_big_object_delayed)(i) for i in range(150)
)
res = accumulator_sum(res)
print('All tasks completed and reduced successfully.')

del res 
monitor_delayed_gen_unordered.join()
peak = max(monitor_delayed_gen_unordered.memory_buffer) / 1e6
print(f"Peak memory usage: {peak:.2f}MB")

Create result generator on delayed tasks with return_as='generator_unordered'...
All tasks completed and reduced successfully.
Peak memory usage: 175.22MB

記憶體使用結果對比如下。基於generator_unordered選項在執行任務時，能夠獨立地處理每個任務，而不需要依賴於其他任務的完成狀態。但是要注意的是由於系統負載、後端實現等多種可能影響任務執行順序的因素，結果的返回順序是不確定的：

plt.figure(1)
plt.semilogy(
    np.maximum.accumulate(monitor_delayed_gen.memory_buffer),
    label='return_as="generator"'
)
plt.semilogy(
    np.maximum.accumulate(monitor_delayed_gen_unordered.memory_buffer),
    label='return_as="generator_unordered"'
)
plt.xlabel("Time")
plt.xticks([], [])
plt.ylabel("Memory usage")
plt.yticks([1e7, 1e8, 1e9], ['10MB', '100MB', '1GB'])
plt.legend()
plt.show()

3 參考

• joblib
• joblib-doc
• loky
• parallel_config
• Python資料序列化模組pickle使用筆記