Python 多核並行計算

以前寫點小程序其實根本不在乎並行，單核跑跑也沒什麼問題，而且我的電腦也只有雙核四個超線程（下面就統稱核好了），覺得去折騰並行沒啥意義（除非在做IO密集型任務）。然後自從用上了32核128GB內存，看到 htop 裡面一堆空載的核，很自然地就會想這個並行必須去折騰一下。後面發現，其實 Python 的並行真的非常簡單。

multiprocessing vs threading

Python 自帶的庫又全又好用，這是我特別喜歡 Python 的原因之一。Python 裡面有 multiprocessing 和 threading 這兩個用來實現並行的庫。用線程應該是很自然的想法，畢竟（直覺上）開銷小，還有共享內存的福利，而且在其他語言裡面線程用的確實是非常頻繁。然而，我可以很負責任的說，如果你用的是 CPython 實現，那麼用了 threading 就等同於和並行計算說再見了（實際上，甚至會比單線程更慢），除非這是個IO密集型的任務。

GIL

CPython 指的是 python.org 提供的 Python 實現。是的，Python 是一門語言，它有各種不同的實現，比如 PyPy, Jython, IronPython 等等……我們用的最多的就是 CPython，它幾乎就和 Python 畫上了等號。

CPython 的實現中，使用了 GIL 即全局鎖，來簡化解釋器的實現，使得解釋器每次只執行一個線程中的位元組碼。也就是說，除非是在等待IO操作，否則 CPython 的多線程就是徹底的謊言！

有關 GIL 下面兩個資料寫的挺好的：

• Python的GIL是什麼鬼，多線程性能究竟如何
• Understanding the Python GIL

multiprocessing.Pool

因為 GIL 的緣故 threading 不能用，那麼我們就好好研究研究 multiprocessing。（當然，如果你說你不用 CPython，沒有 GIL 的問題，那也是極佳的。）

首先介紹一個簡單粗暴，非常實用的工具，就是 multiprocessing.Pool。如果你的任務能用 ys = map(f, xs) 來解決，大家可能都知道，這樣的形式天生就是最容易並行的，那麼在 Python 裡面並行計算這個任務真是再簡單不過了。舉個例子，把每個數都平方：

import multiprocessingdef f(x): return x * xcores = multiprocessing.cpu_count()pool = multiprocessing.Pool(processes=cores)xs = range(5)# method 1: mapprint pool.map(f, xs) # prints [0, 1, 4, 9, 16]# method 2: imapfor y in pool.imap(f, xs): print y # 0, 1, 4, 9, 16, respectively# method 3: imap_unorderedfor y in pool.imap_unordered(f, xs): print(y) # may be in any order

map 直接返回列表，而 i 開頭的兩個函數返回的是迭代器；imap_unordered 返回的是無序的。

當計算時間比較長的時候，我們可能想要加上一個進度條，這個時候 i 系列的好處就體現出來了。另外，有一個小技巧，就是輸出
可以使得游標回到行首而不換行，這樣就可以製作簡易的進度條了。

cnt = 0for _ in pool.imap_unordered(f, xs): sys.stdout.write(done %d/%d
% (cnt, len(xs))) cnt += 1

更複雜的操作

要進行更複雜的操作，可以直接使用 multiprocessing.Process 對象。要在進程間通信可以使用：

• multiprocessing.Pipe
• multiprocessing.Queue
• 同步原語
• 共享變數

其中我強烈推薦的就是 Queue，因為其實很多場景就是生產者消費者模型，這個時候用 Queue 就解決問題了。用的方法也很簡單，現在父進程創建 Queue，然後把它當做 args 或者 kwargs 傳給 Process 就好了。

使用 Theano 或者 Tensorflow 等工具時的注意事項

需要注意的是，在 import theano 或者 import tensorflow 等調用了 Cuda 的工具的時候會產生一些副作用，這些副作用會原樣拷貝到子進程中，然後就發生錯誤，如：

could not retrieve CUDA device count: CUDA_ERROR_NOT_INITIALIZED

解決的方法是，保證父進程不引入這些工具，而是在子進程創建好了以後，讓子進程各自引入。

如果使用 Process，那就在 target 函數裡面 import。舉個例子：

import multiprocessingdef hello(taskq, resultq): import tensorflow as tf config = tf.ConfigProto() config.gpu_options.allow_growth=True sess = tf.Session(config=config) while True: name = taskq.get() res = sess.run(tf.constant(hello + name)) resultq.put(res)if __name__ == __main__: taskq = multiprocessing.Queue() resultq = multiprocessing.Queue() p = multiprocessing.Process(target=hello, args=(taskq, resultq)) p.start() taskq.put(world) taskq.put(abcdabcd987) taskq.close() print(resultq.get()) print(resultq.get()) p.terminate() p.join()

如果使用 Pool，那麼可以編寫一個函數，在這個函數裡面 import，並且把這個函數作為 initializer 傳入到 Pool 的構造函數裡面。舉個例子：

import multiprocessingdef init(): global tf global sess import tensorflow as tf config = tf.ConfigProto() config.gpu_options.allow_growth=True sess = tf.Session(config=config)def hello(name): return sess.run(tf.constant(hello + name))if __name__ == __main__: pool = multiprocessing.Pool(processes=2, initializer=init) xs = [world, abcdabcd987, Lequn Chen] print pool.map(hello, xs)