首頁?學習  »   正文

不用多進程的Python十倍速并行技巧(下)

上一篇我們學習了三種不易用Python多處理表示的工作負載基準測試的其中兩種,并比較了Ray、Python多處理和串行Python代碼。今天這一篇我們來聊聊第三種基準測試。

傳送門:不用多進程的Python十倍速并行技巧(上)

基準3:昂貴的初始化

與前面的示例不同,許多并行計算不一定要求在任務之間共享中間計算,但無論如何都會從中受益。即使是無狀態計算,在狀態初始化代價高昂時也可以從共享狀態中獲益。

下面是一個例子,我們希望從磁盤加載一個保存的神經網絡,并使用它來并行分類一組圖像。

不用多進程的Python十倍速并行技巧(下)

在48個物理內核的機器上,Ray比Python多處理快25倍,比單線程Python快13倍。在本例中,Rython多處理并沒有優于單線程Rython。錯誤欄被描繪出來,但在某些情況下太小而看不見。工作負載被縮放到核心的數量,因此在更多的核心上完成更多的工作。在這個基準測試中,“串行”Python代碼實際上通過TensorFlow使用多個線程。Python多處理代碼的可變性來自于重復從磁盤加載模型的可變性,而其他方法不需要這樣做。

本例使用Ray的5s、Python多處理的126s和串行Python的64s(在48個物理內核上)。在本例中,串行Python版本使用多個內核(通過TensorFlow)并行化計算,因此它實際上不是單線程的。

假設我們最初通過運行以下代碼創建了模型:

import tensorflow as tf

mnist = tf.keras.datasets.mnist.load_data()
x_train, y_train = mnist[0]
x_train = x_train / 255.0
model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(512, activation=tf.nn.relu),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(10, activation=tf.nn.softmax)
])
model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy'])
# Train the model.
model.fit(x_train, y_train, epochs=1)
# Save the model to disk.
filename = '/tmp/model'
model.save(filename)

用于將神經網絡模型保存到磁盤的代碼。

現在,我們希望加載模型并使用它對一組圖像進行分類。我們批量進行這項工作是因為在應用程序中,圖像可能不會全部同時可用,而圖像分類可能需要與數據加載并行進行。

Ray版本如下所示:

import psutil
import ray
import sys
import tensorflow as tf

num_cpus = psutil.cpu_count(logical=False)

ray.init(num_cpus=num_cpus)

filename = '/tmp/model'

@ray.remote
class Model(object):
    def __init__(self, i):
        # Pin the actor to a specific core if we are on Linux to prevent
        # contention between the different actors since TensorFlow uses
        # multiple threads.
        if sys.platform == 'linux':
            psutil.Process().cpu_affinity([i])
        # Load the model and some data.
        self.model = tf.keras.models.load_model(filename)
        mnist = tf.keras.datasets.mnist.load_data()
        self.x_test = mnist[1][0] / 255.0

    def evaluate_next_batch(self):
        # Note that we reuse the same data over and over, but in a
        # real application, the data would be different each time.
        return self.model.predict(self.x_test)

actors = [Model.remote(i) for i in range(num_cpus)]

# Time the code below.

# Parallelize the evaluation of some test data.
for j in range(10):
    results = ray.get([actor.evaluate_next_batch.remote() for actor in actors])

使用Ray的玩具分類示例的代碼。

加載模型的速度很慢,我們只想加載一次。Ray版本通過在參與者的構造函數中加載模型一次來分攤此成本。如果模型需要放在GPU上,那么初始化將更加昂貴。

多處理版本速度較慢,因為它需要在每次映射調用中重新加載模型,因為假定映射函數是無狀態的。

多處理版本如下。注意,在某些情況下,可以使用multiprocessing.Pool的初始化參數來實現這一點。但是,這僅限于初始化對每個進程都是相同的設置,并且不允許不同的進程執行不同的設置函數(例如,加載不同的神經網絡模型),也不允許針對不同的工作者執行不同的任務。

from multiprocessing import Pool
import psutil
import sys
import tensorflow as tf

num_cpus = psutil.cpu_count(logical=False)

filename = '/tmp/model'

def evaluate_next_batch(i):
    # Pin the process to a specific core if we are on Linux to prevent
    # contention between the different processes since TensorFlow uses
    # multiple threads.
    if sys.platform == 'linux':
        psutil.Process().cpu_affinity([i])
    model = tf.keras.models.load_model(filename)
    mnist = tf.keras.datasets.mnist.load_data()
    x_test = mnist[1][0] / 255.0
    return model.predict(x_test)

pool = Pool(num_cpus)

for _ in range(10):
    pool.map(evaluate_next_batch, range(num_cpus))

使用多處理的玩具分類示例的代碼。

我們在所有這些例子中看到的是,Ray的性能不僅來自于它的性能優化,還來自于對手頭任務進行適當的抽象。有狀態計算對許多應用程序都很重要,將有狀態計算強制為無狀態抽象是有代價的。

運行基準測試

在運行這些基準測試之前,您需要安裝以下軟件。

pip install numpy psutil ray scipy tensorflow

然后,通過運行這些腳本,可以復制上面的所有數字。

如果在安裝psutil時遇到問題,請嘗試使用Python。

最初的基準測試是使用M5實例類型(M5.large用于1個物理內核,M5.24XLarge用于48個物理內核)在EC2上運行的。

為了使用正確的配置啟動AWS或GCP上的實例,您可以使用Ray Autoscaler并運行以下命令。

ray up config.yaml

這里提供了一個config.yaml示例(用于啟動M5.4XL實例)。

關于Ray的更多信息

盡管本文關注的是Ray和Python多處理之間的基準測試,但蘋果對蘋果的比較是具有挑戰性的,因為這些庫并不十分相似。差異包括以下內容。

Ray是為可擴展性而設計的,可以在筆記本電腦和集群上運行相同的代碼(多處理僅在一臺機器上運行)。

Ray工作負載自動從機器和流程故障中恢復。

Ray是以一種與語言無關的方式設計的,并且對Java有初步的支持。

所以說,這篇文章幫到你了嗎?謝謝捧場~

歡迎關注ATYUN官方公眾號,商務合作及內容投稿請聯系郵箱:[email protected]

發表評論