原文地址：http://yongyuan.name/blog/caffe-python-extract-feature-parallel.html

https://thief.one/2016/11/24/Multiprocessing-Pool/

https://jingsam.github.io/2015/12/31/multiprocessing.html

想象這樣一個場景：伺服器上配備了很多GPU顯示卡，而你又使用的是Caffe，不幸的是，你還選用了Python來寫程式碼，不幸中的不幸是你還要在短時間內處理百萬千萬規模的圖片。那麼，問題來了，Caffe本身對多卡的支援就不是很好，而你還要用Python，而且即便你通過設定batch size的大小來加快處理速度，但你還是隻把一張顯示卡用起來。有沒有辦法把所有的GPU都用起來，並行提取特徵呢？

上面這個問題在實際中是一個經常會碰到並且繞不過的問題，小白菜在近期也遇到了此問題，並且上面那些設定的條件小白菜不幸都一一躺槍。雖然如此，問題還是要解決的。下面小白菜聊聊如何使用python 的multiprocessing多程序並行庫來把伺服器上的多張GPU都用起來，進行特徵的並行提取。在正式給出具體的解決方案之前，還是有必要了解multiprocessing多程序並行庫的。

multiprocessing多程序並行庫

multiprocessing對於程序的管理，是使用程序池pool進行管理的。pool允許使用者指定程序數目，線上程池未滿的時候，如果新的請求提交過來，則pool會建立新的程序來執行該請求，如果池中的程序數已經達到規定最大值，那麼該請求就會等待，直到池中有程序結束，才會建立新的程序來執行它。pool對於程序的管理，有兩種方式，分別是非同步方式(也稱非阻塞方式)和同步方式(也稱阻塞方式)，對應的方法為apply_async()和apply()。

• apply(): 同步方式(也稱阻塞方式)。阻塞，顧名思義，在執行該任務的時候，把所有的介面資源都堵住了，其他的任務必須等待這個任務執行完後，再釋放資源後，才能執行其他的任務，所以非常的好理解。
• pool.apply_async(): 非同步方式(也稱非阻塞方式)。有了上面阻塞的解讀，小白菜以為非同步方式已經非常直白了。

程序池管理pool還有一些其他的方法，這裡不再做詳細的展開了，更多關於Python的multiprocessing的解讀，可以參考Multiprocessing.Pool和正確使用 Multiprocessing 的姿勢，寫得直白易懂。

有了上面對Python的multiprocessing的認識，我們再回到使用python的multiprocessing多程序並行庫來把伺服器上的多張GPU都用起來這個問題上。根據對pool對程序池管理的方法，顯然pool.apply_async()

並行化方案

問題定義：假設需要處理的圖片有800萬，伺服器上有8塊GPU，如何使用python 的multiprocessing多程序並行庫來把伺服器上的多張GPU都用起來，進行特徵的並行提取。

對於上述給出的問題定義，小白菜設計的並行化方案為：將800萬分成8等分，啟動8個程序，每個程序使用一塊GPU，並且每一個程序載入一份預訓練的模型，這樣就可以把8塊GPU都用起來了。此種方案，可以最大限度的使用GPU，有多少塊GPU，就可以用多少GPU，下面是具體的編碼。

具體編碼

有了上面對multiprocessing和並行化方案後，在具體編碼的時候，要做的事情其實非常簡單，主要需要完成的就是兩件事：

• 把圖片資料按GPU的數目進行等量劃分，即寫一個split_list()函式
• 正常編寫特徵提取程式碼，寫一個gpu_task函式

具體程式碼如下:

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Author: [email protected]


import multiprocessing
from multiprocessing import Process, freeze_support, Pool
import sys
import os
import caffe
import numpy as np
import scipy
import h5py


def split_list(alist, wanted_parts=1):
    length = len(alist)
    return [ alist[i*length // wanted_parts: (i+1)*length // wanted_parts] 
             for i in range(wanted_parts) ]

def gpu_task(prototxt, caffemodel, layer, path_images, out, gpu=0):
    num_images = len(path_images)

    h5f = h5py.File(out, 'w')

    # set gpu card
    caffe.set_device(gpu)
    caffe.set_mode_gpu()

    # init NN
    net = caffe.Net(prototxt, caffemodel, caffe.TEST)
    net.forward()

    features = []
    image_names = []

    for i, path in enumerate(path_images):
        print "%d(%d), %s"%((i+1), num_images, os.path.basename(path))
        if os.path.splitext(os.path.basename(path))[0][-6:] != 'single':
            img = ...
 
            # extract feature

            feat = ...
 
            features.append(feat)
            image_names.append(os.path.basename(path))      

    features = np.array(features)
    h5f['feats'] = features
    h5f['names'] = image_names
    h5f.close()    
    print "gpu %d task has finished..." % (gpu)

if __name__ == '__main__':

    multiprocessing.freeze_support()
    pool = multiprocessing.Pool()
    
    dir_images = '....'
    path_images = [os.path.join(dir_images, f) for f in sorted(os.listdir(dir_images)) if f.endswith('.jpg')]
    layer = '...'
    prototxt = 'VGG_ILSVRC_16.prototxt'
    caffemodel = 'VGG_ILSVRC_16_layers.caffemodel'
    out = '...'

    parts = 8 # 8個程序

    out_files = []
    for i in xrange(parts):
        out_files.append(os.path.join(out, str(i) + '.h5'))

    blocks = split_list(path_images, wanted_parts = parts)
  
    for i in xrange(0, parts):
        pool.apply_async(gpu_task, args = (prototxt, caffemodel, layer, blocks[i], out_files[i], i + (8 - parts),))
    
    pool.close()
    pool.join()

上面那些...的部分，是小白菜刻意省略的，具體在實際使用的時候，可以結合自己的需要來做更改。另外一個pool.close()和pool.join()的次序一定不能反，否則會出錯。執行完close後不會有新的程序加入到pool,join函式等待所有子程序結束。

對於特徵的儲存，可以採用hdf5的格式進行儲存，上面每一個程序提取的特徵，儲存在一個h5(HDF5格式)檔案中，這樣對於八個程序，可以得到8份小檔案，每一份檔案裡存了影象檔名及其對應的特徵。