深度解決添加復(fù)雜數(shù)據(jù)增強(qiáng)導(dǎo)致訓(xùn)練模型耗時長的痛點(diǎn)（1）

本文來自社區(qū)投稿

作者：教 主

原文鏈接：

https://zhuanlan.zhihu.com/p/585270139

作者薦語：

最近在訓(xùn)練大規(guī)模數(shù)據(jù)時，遇到一個【添加復(fù)雜數(shù)據(jù)增強(qiáng)導(dǎo)致訓(xùn)練模型耗時長】的問題，在學(xué)習(xí)了 MMDetection 和 MMCV 底層關(guān)于 PyTorch 的 CUDA/C++ 拓展之后，我也將一些復(fù)雜數(shù)據(jù)增強(qiáng)實(shí)現(xiàn)了 GPU 化，并且詳細(xì)總結(jié)了一些經(jīng)驗(yàn)，分享此篇文章和工程，希望與大家多多交流。

同時感謝 MMDetection 和 MMCV 提供的寶貴 code。

0 Introduction

一直以來，得益于 GPU 的日益發(fā)展，深度學(xué)習(xí)中網(wǎng)絡(luò)訓(xùn)練以及部署推理速度越來越快，在各大主流的深度學(xué)習(xí)框架，諸如 PyTorch、TensorFlow、OneFlow 等都有很多算子對 GPU 的加速支持。

從網(wǎng)絡(luò)結(jié)構(gòu)角度，PyTorch 雖然已經(jīng)使用了 NVIDIA cuDNN、Intel MKL 和 NNPACK 這些底層來加快訓(xùn)練速度，但是在某些情況下，比如我們要實(shí)現(xiàn)一些特定算法/算子，如果只是用 PyTorch 已有的算子或操作遠(yuǎn)遠(yuǎn)不夠。

因?yàn)?PyTorch 雖然在特定操作上經(jīng)過了很好的優(yōu)化，但是對于 PyTorch 已經(jīng)寫好的這些操作，假如我們組合起來成為一個新的算子（OP），PyTorch 不會管你的算法的具體執(zhí)行流程，一般 PyTorch 只會按照設(shè)計好的操作去使用 GPU，然后 GPU 可能不能充分利用或者直接超負(fù)載，并且 python 解釋器也不能對此進(jìn)行優(yōu)化，導(dǎo)致訓(xùn)練過程變慢很多 [1]。

從數(shù)據(jù)流角度，深度學(xué)習(xí)一般都需要復(fù)雜的、多階段的數(shù)據(jù)處理流程，包括數(shù)據(jù)加載、解碼以及一定量的數(shù)據(jù)增強(qiáng)預(yù)處理操作，這些目前在 CPU 上執(zhí)行的數(shù)據(jù)處理管道已經(jīng)成為瓶頸，使得模型訓(xùn)練耗時很長大。

對于此，NVIDIA 提出了 Data Loading Library（DALI）[2]，通過將數(shù)據(jù)預(yù)處理交給 GPU 處理，緩解 CPU 瓶頸問題。DALI 依賴于它自己的執(zhí)行引擎，其構(gòu)建目的是最大化輸入管道的吞吐量。諸如預(yù)取、并行執(zhí)行和批處理等特性都是為用戶透明處理，如下圖所示：

DALI Pipeline

使用 DALI 以及配置 DALI 環(huán)境比較復(fù)雜，并且 DALI 當(dāng)前的支持的函數(shù)實(shí)現(xiàn)也比較有限，具體使用可以看文獻(xiàn) [2] 中的說明文檔。

實(shí)際開發(fā)中，對于一些復(fù)雜的特定數(shù)據(jù)增強(qiáng)操作，就需要自己實(shí)現(xiàn)。因此，構(gòu)建了一個比較全面的工程以供大家學(xué)習(xí)和相互交流。

本工程利用 Pytorch 的 C++/CUDA 擴(kuò)展，實(shí)現(xiàn) GPU 的數(shù)據(jù)增強(qiáng)，然后直接推送給網(wǎng)絡(luò)，從而達(dá)到訓(xùn)練加速效果。

為了指導(dǎo)大家系統(tǒng)性掌握該方面的相關(guān)知識，本工程也包含了 Python 的 C++ 拓展，且詳細(xì)講解了在需要依賴第三方庫的情況下怎樣編寫 setup.py 文件以及相關(guān)配置，關(guān)于如何編譯和測試，在后續(xù)有詳細(xì)的講解。

1. Project Address

https://github.com/ChenCVer/python_cpp_extension

2. Project Structure

├── 3rdparty          # 工程依賴的第三方庫│    ├── opencv│    │    ├── linux│    │    └── win│    └── pybind11├── docs              # 說明文檔及相關(guān)資料├── requirements      # python相關(guān)安裝依賴├── requirements.txt  # python相關(guān)安裝依賴項(xiàng), 與requirements文件夾配合├── scripts           # 相關(guān)測試腳本├── tools             # 分析工具├── orbbec            # 源碼文件│    ├── nms          # 非極大值抑制│    ├── roi_align    # ROI Align│    ├── utils        # 編譯工具函數(shù)│    └── warpaffine   # 仿射變換增強(qiáng)└── setup.py          # 用于編譯和構(gòu)建python包(.egg), 類似:CMakeLists.txt

3. Compilation And Python Environment

3.1. Compile Environment

• GCC/G++ >= 5.5.0(Visual Studio 2017 or newer for Windows)

• CUDA(NVCC): 10.1~11.5

3.2. Python Environment

（requirements.txt）

certifi==2021.5.30cycler==0.11.0future==0.18.2kiwisolver==1.3.1matplotlib==3.3.4mkl-fft==1.3.0mkl-random==1.1.1mkl-service==2.3.0numpy @ file:///C:/ci/numpy_and_numpy_base_1603480701039/workolefile==0.46opencv-python==3.4.0.12Pillow @ file:///C:/ci/pillow_1625663293114/workpyparsing==3.0.9python-dateutil==2.8.2six @ file:///tmp/build/80754af9/six_1644875935023/workterminaltables==3.1.10torch==1.5.0torchvision==0.6.0wincertstore==0.2

3.3. Python Package infos

Package         Version --------------- --------- certifi         2016.2.28cycler          0.11.0Cython          0.29.32future          0.18.2kiwisolver      1.3.1matplotlib      3.3.4mkl-fft         1.3.0mkl-random      1.1.1mkl-service     2.3.0numpy           1.19.2olefile         0.44opencv-python   3.4.0.12Pillow          8.3.1pip             21.3.1pyparsing       3.0.9python-dateutil 2.8.2setuptools      59.6.0six             1.10.0terminaltables  3.1.10torch           1.5.0torchvision     0.6.0wheel           0.29.0wincertstore    0.2

【注】：上述環(huán)境中的 PyTorch 版本需要對應(yīng)的 CUDA 版本，本工程支持的 PyTorch 版本：PyTorch version：1.5.0~latest。

4. C++ And CUDA Extensions 

For Python/ PyTorch

C++ 與 Python 或 PyTorch 的交互，業(yè)界主流做法是采用 pybind11，關(guān)于Pybind11 的更多詳細(xì)說明可以參看文獻(xiàn) [15]，其核心原理如下圖所示：

pybind11 pipeline

由于 PyTorch 的 C++ 拓展與純 Python 有一些區(qū)別，因?yàn)?PyTorch 的基礎(chǔ)數(shù)據(jù)類型是 torch.Tensor，該數(shù)據(jù)類型可以認(rèn)為是 Pytorch 庫對 np.array 進(jìn)行了更高一層的封裝。所以，在寫拓展程序時，其接口函數(shù)所需要的數(shù)據(jù)類型以及調(diào)用的庫會有些區(qū)別，下面會詳細(xì)解釋。

4.1. C++ Extensions For Python

首先我們看 Python 代碼，如下所示（scripts/test_warpaffine_opencv.py）：

import cv2import torch  # 不能刪掉, 因?yàn)樾枰獎討B(tài)加載torch的一些動態(tài)庫,后面會詳細(xì)說明.import numpy as npfrom orbbec.warpaffine import affine_opencv  # C++ interface
data_path = "./demo.png"img = cv2.imread(data_path, cv2.IMREAD_GRAYSCALE)
# python中的numpy.array()與 pybind中的py::array_t一一對應(yīng).src_point = np.array([[262.0, 324.0], [325.0, 323.0], [295.0, 349.0]], dtype=np.float32)dst_point = np.array([[38.29, 51.69], [73.53, 51.69], [56.02, 71.73]], dtype=np.float32)# python interface mat_trans = cv2.getAffineTransform(src_point, dst_point)res = cv2.warpAffine(img, mat_trans, (600,800))cv2.imwrite("py_img.png", res)
# C++ interfacewarpffine_img = affine_opencv(img, src_point, dst_point)cv2.imwrite("cpp_img.png", warpffine_img)

從上述代碼可以看到，Python 文件中調(diào)用了 affine_opencv 函數(shù)，而 affine_opencv 的 C++ 實(shí)現(xiàn)在 orbbec/warpaffine/src/cpu/warpaffine_opencv.cpp 中，如下所示：

#include<vector>#include<iostream>#include<pybind11/pybind11.h>#include<pybind11/numpy.h>#include<pybind11/stl.h>#include<opencv2/opencv.hpp>

namespace py = pybind11;
/* Python->C++ Mat */cv::Mat numpy_uint8_1c_to_cv_mat(py::array_t<unsigned char>& input){        ...}
cv::Mat numpy_uint8_3c_to_cv_mat(py::array_t<unsigned char>& input){        ...}
/* C++ Mat ->numpy */py::array_t<unsigned char> cv_mat_uint8_1c_to_numpy(cv::Mat& input){        ...}
py::array_t<unsigned char> cv_mat_uint8_3c_to_numpy(cv::Mat& input){        ...}
py::array_t<unsigned char> affine_opencv(py::array_t<unsigned char>& input,                                         py::array_t<float>& from_point,                                         py::array_t<float>& to_point){        ...}

由于本工程同時兼容了 PyTorch 的 C++/CUDA 拓展，為了更加規(guī)范，這里在拓展接口程序（orbbec/warpaffine/src/warpaffine_ext.cpp）中通過 PYBIND11_MODULE 定義好接口，如下所示：

#include <torch/extension.h>#include<pybind11/numpy.h>
// python的C++拓展函數(shù)申明py::array_t<unsigned char> affine_opencv(py::array_t<unsigned char>& input,                                         py::array_t<float>& from_point,                                         py::array_t<float>& to_point);
// Pytorch的C++拓展函數(shù)申明(CPU)at::Tensor affine_cpu(const at::Tensor& input,          /*[B, C, H, W]*/                      const at::Tensor& affine_matrix,  /*[B, 2, 3]*/                      const int out_h,                      const int out_w);
// Pytorch的CUDA拓展函數(shù)申明(GPU)#ifdef WITH_CUDAat::Tensor affine_gpu(const at::Tensor& input,          /*[B, C, H, W]*/                      const at::Tensor& affine_matrix,  /*[B, 2, 3]*/                      const int out_h,                      const int out_w);#endif
// 通過WITH_CUDA宏進(jìn)一步封裝Pytorch的拓展接口at::Tensor affine_torch(const at::Tensor& input,          /*[B, C, H, W]*/                                  const at::Tensor& affine_matrix,  /*[B, 2, 3]*/                                  const int out_h,                                  const int out_w){        if (input.device().is_cuda())          {#ifdef WITH_CUDA    return affine_gpu(input, affine_matrix, out_h, out_w);#else    AT_ERROR("affine is not compiled with GPU support");#endif          }          return affine_cpu(input, affine_matrix, out_h, out_w);}
// 使用pybind11模塊定義python/pytorch接口PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {  m.def("affine_opencv", &affine_opencv, "affine with c++ opencv");  m.def("affine_torch", &affine_torch,   "affine with c++ libtorch");}

從上面代碼可以看出，Python 中的 np.array 數(shù)組與 pybind11 的 py::array_t 相互對應(yīng)，也即 Python 接口函數(shù)中，傳入的 np.array 數(shù)組，在 C++ 對應(yīng)的函數(shù)中用 py::array_t 接收，操作 Numpy 數(shù)組，需要引入頭文件。

數(shù)組本質(zhì)上在底層是一塊一維的連續(xù)內(nèi)存區(qū)，通過 pybind11 中的 request() 函數(shù)可以把數(shù)組解析成 py::buffer_info 結(jié)構(gòu)體，buffer_info 類型可以公開一個緩沖區(qū)視圖，它提供對內(nèi)部數(shù)據(jù)的快速直接訪問，如下代碼所示：

struct buffer_info {    void *ptr;                         // 指向數(shù)組(緩沖區(qū))數(shù)據(jù)的指針    py::ssize_t itemsize;              // 數(shù)組元素總數(shù)    std::string format;                // 數(shù)組元素格式(python表示的類型)    py::ssize_t ndim;                  // 數(shù)組維度信息    std::vector<py::ssize_t> shape;    // 數(shù)組形狀    std::vector<py::ssize_t> strides;  // 每個維度相鄰元素的間隔(字節(jié)數(shù)表示)};

在寫好 C++ 源碼以后，在 setup.py 中將相關(guān) C++ 源文件，以及依賴的第三方庫：opencv、pybind11 的路徑寫入對應(yīng)位置（本工程已經(jīng)寫好，請具體看 setup.py 文件），然后進(jìn)行編譯和安裝：

# 切換工作路徑step 1: cd F:/code/python_cpp_extension# 編譯step 2: python setup.py develop# 安裝, 如果沒有指定--prefix, 則最終編譯成功的安裝包(.egg)文件會安裝到對應(yīng)的python環(huán)境下的site-packages下.step 3: python setup.py install

【注】：關(guān)于工程文件中的 setup.py 相關(guān)知識可以參考文獻(xiàn) [7]、[12]、[13]，該三篇文獻(xiàn)對此有詳細(xì)的解釋。

執(zhí)行 step 2 和 step3 之后，如下圖所示，最終源碼文件會編譯成 .pyd 二進(jìn)制文件（Linux 系統(tǒng)下編譯成 .so 文件），且會生成一個 Python 包文件：orbbec-0.0.1-py36-win-amd64.egg，包名取決于 setup.py 中規(guī)定的 name 和 version 信息，該安裝包會被安裝在當(dāng)前 Python環(huán)境的 site-packages 文件夾下。

同時，在終端執(zhí)行命令：pip list，會發(fā)現(xiàn)安裝包以及對應(yīng)的版本信息。安裝成功后，也就意味著，在該 Python環(huán)境（本工程的 Python環(huán)境是 cpp_extension）下，可以在任何一個 Python 文件中，導(dǎo)入 orbbec 安裝包中的接口函數(shù)，比如上述 scripts/test_warpaffine_opencv.py 文件中的語句：from orbbec.warpaffine import affine_opencv。

編譯和安裝成功

pip list 顯示相關(guān)安裝包信息