想象這樣一個場景：服務器上配備了很多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: github@willard-yuanimport multiprocessing from multiprocessing import Process, freeze_support, Pool import sys import os import caffe import numpy as np import scipy import h5pydef 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 cardcaffe.set_device(gpu)caffe.set_mode_gpu()# init NNnet = 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 featurefeat = ...features.append(feat)image_names.append(os.path.basename(path)) features = np.array(features)h5f['feats'] = featuresh5f['names'] = image_namesh5f.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份小文件，每一份文件里存了圖像文件名及其對應的特征。

from:http://yongyuan.name/blog/caffe-python-extract-feature-parallel.html?

總結

以上是生活随笔為你收集整理的机器视觉：Caffe Python接口多进程提取特征的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。