pytorch 網(wǎng)絡(luò)預(yù)處理與后處理中基于numpy操作的GPU加速

　　背景

　　python腳本運行在服務(wù)器端的卷積神經(jīng)網(wǎng)絡(luò)往往需要將圖片數(shù)據(jù)從cv2(numpy.ndarray)->tensor送入網(wǎng)絡(luò)，之后進行inference，再將結(jié)果從tensor-> numpy.ndarray的過程。

　　由于cv2讀取的數(shù)據(jù)存于內(nèi)存中，以pytorch框架舉例，在把數(shù)據(jù)送入GPU前會產(chǎn)生如下的數(shù)據(jù)轉(zhuǎn)換：

　　GPU準備進行inference之前會判斷torch.cuda.FloatTensor是否已經(jīng)處于顯存內(nèi)，如果沒有的話會隱式調(diào)用內(nèi)存與顯存中的數(shù)據(jù)轉(zhuǎn)存協(xié)議方法.async_copy()函數(shù)，將數(shù)據(jù)轉(zhuǎn)存至GPU顯存中，但該部分往往需要消耗大量時間。

　　對策：直接在GPU顯存中開辟空間

　　應(yīng)用庫：cupy、dlpack

　　一、前處理

　　通常pytorch前處理如下：

　　# 內(nèi)存分配torch.FloatTensor空間

　　batch_input = torch.zeros(len(image_list), 3, target_height, target_width)

　　for index in range(len(image_list)):

　　# image->numpy.ndarray

　　img = cv2.resize(image_list[index].copy(), (target_width, target_height))

　　# uint8->float32

　　t_img = np.asarray(img, np.float32)

　　#轉(zhuǎn)置

　　m_img = t_img.transpose((2, 0, 1))

　　#numpy.ndarray->torch.FloatTensor + 圖像正則化

　　n_img = transform(torch.from_numpy(m_img))

　　#組成batch data

　　batch_input[index, :] = n_img

　　# torch.FloatTensor-> torch.cuda.FloatTensor

　　batch_input.cuda()

　　如果將此batch送入GPU，則會發(fā)生如圖1所示的數(shù)據(jù)轉(zhuǎn)換。

　　現(xiàn)用cupy來取代numpy操作:

　　import cupy as cp

　　# GPU顯存分配cupy batch_data空間

　　batch_input = cp.zeros((len(image_list), 3, target_height, target_width), dtype=cp.float32)

　　for index in range(len(image_list)):

　　# image->cupy.ndarray

　　img = cv2.resize(image_list[index], (target_width, target_height))

　　# numpy.uint8 -> cupy.float32

　　t_img = cp.asarray(img, cp.float32)

　　# 轉(zhuǎn)置(cupy層面)

　　m_img = t_img.transpose((2, 0, 1))

　　# 圖像正則化

　　n_img = gpu_transform(m_img)

　　# 組成 batch data

　　batch_input[index, :] = n_img

　　# cupy.ndarray -> torch.cuda.FloatTensor

　　batch_data = from_dlpack(toDlpack(batch_input)).cuda()

　　此時過程轉(zhuǎn)換為：

　　說明幾點：

　　1.1由于cupy直接在GPU顯存中分配空間，不需要隱式調(diào)用.async_copy()將數(shù)據(jù)調(diào)入顯存內(nèi)，可見時間對比：

　　隱式調(diào)用GPU前傳時間如下圖：

　　非隱式調(diào)用GPU前傳時間如下圖：

　　1.2 cupy.ndarray到torch.cuda.FloatTensor沒辦法直接轉(zhuǎn)換，需要中間轉(zhuǎn)換格式dlpack，具體轉(zhuǎn)換如下

　　rom cupy.core.dlpack import toDlpack

　　from cupy.core.dlpack import fromDlpack

　　from torch.utils.dlpack import to_dlpack

　　from torch.utils.dlpack import from_dlpack

　　import torch鄭州婦科醫(yī)院 http://www.sptdfk.com/

　　#tensor->cupy

　　cupy_data = fromDlpack(to_dlpack(tensor_data))

　　#cupy->tensor

　　tensor_data = from_dlpack(toDlpack(cupy_data))

　　1.3 在pytorch框架中，有的工程需要圖像正則化，有的不需要。當網(wǎng)絡(luò)前傳時若需要圖像正則化(一般為減均值與除方差)，一般選用的是torchvision.transform。但是該內(nèi)置函數(shù)只接受CPU端的torch.FloatTensor，這就意味著若要使用內(nèi)置transform函數(shù)，就需要將cupy GPU數(shù)據(jù)先轉(zhuǎn)成CPU的torch.FloatTensor，勢必會造成數(shù)據(jù)轉(zhuǎn)換資源浪費。重寫transform函數(shù)：

　　self.mean = cp.array([102.9801, 115.9465, 122.7717])

　　self.std = cp.array([1., 1., 1.])

　　def gpu_transform(self, img):

　　for index in range(img.shape[0]):

　　img[index,:] -= self.mean[index]

　　img[index, :] /= self.std[index]

　　return img

　　以上過程全部都在GPU內(nèi)運行，時間幾乎可以忽略

　　二、后處理

　　此部分適用于分割網(wǎng)絡(luò)，即需要預(yù)先在GPU端分配生成的mask空間。通常做法分配torch.cuda.FloatTensor空間，隱式調(diào)用.async_copy()送入GPU，同樣會消耗很多時間。類似于前處理，可以利用cupy生成mask空間，再轉(zhuǎn)torch.cuda.FloatTensor。

　　mask_gpu= from_dlpack(toDlpack(cp.zeros((len(image_list), self.num_classes, ori_img_size[0], ori_img_size[1]), dtype=cp.float32))).cuda()

　　pytorch分配mask時間

　　cupy分配mask時間

　　三、cupy與常規(guī)前后處理時間對比