深度解決添加復雜數據增強導致訓練模型耗時長的痛點（1）

本文來自社區投稿

作者：教 主

原文鏈接：

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

作者薦語：

最近在訓練大規模數據時，遇到一個【添加復雜數據增強導致訓練模型耗時長】的問題，在學習了 MMDetection 和 MMCV 底層關于 PyTorch 的 CUDA/C++ 拓展之后，我也將一些復雜數據增強實現了 GPU 化，并且詳細總結了一些經驗，分享此篇文章和工程，希望與大家多多交流。

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

0 Introduction

一直以來，得益于 GPU 的日益發展，深度學習中網絡訓練以及部署推理速度越來越快，在各大主流的深度學習框架，諸如 PyTorch、TensorFlow、OneFlow 等都有很多算子對 GPU 的加速支持。

從網絡結構角度，PyTorch 雖然已經使用了 NVIDIA cuDNN、Intel MKL 和 NNPACK 這些底層來加快訓練速度，但是在某些情況下，比如我們要實現一些特定算法/算子，如果只是用 PyTorch 已有的算子或操作遠遠不夠。

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

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

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

DALI Pipeline

使用 DALI 以及配置 DALI 環境比較復雜，并且 DALI 當前的支持的函數實現也比較有限，具體使用可以看文獻 [2] 中的說明文檔。

實際開發中，對于一些復雜的特定數據增強操作，就需要自己實現。因此，構建了一個比較全面的工程以供大家學習和相互交流。

本工程利用 Pytorch 的 C++/CUDA 擴展，實現 GPU 的數據增強，然后直接推送給網絡，從而達到訓練加速效果。

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

1. Project Address

https://github.com/ChenCVer/python_cpp_extension

2. Project Structure

├── 3rdparty          # 工程依賴的第三方庫│    ├── opencv│    │    ├── linux│    │    └── win│    └── pybind11├── docs              # 說明文檔及相關資料├── requirements      # python相關安裝依賴├── requirements.txt  # python相關安裝依賴項, 與requirements文件夾配合├── scripts           # 相關測試腳本├── tools             # 分析工具├── orbbec            # 源碼文件│    ├── nms          # 非極大值抑制│    ├── roi_align    # ROI Align│    ├── utils        # 編譯工具函數│    └── warpaffine   # 仿射變換增強└── setup.py          # 用于編譯和構建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

【注】：上述環境中的 PyTorch 版本需要對應的 CUDA 版本，本工程支持的 PyTorch 版本：PyTorch version：1.5.0~latest。

4. C++ And CUDA Extensions 

For Python/ PyTorch

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

pybind11 pipeline

由于 PyTorch 的 C++ 拓展與純 Python 有一些區別，因為 PyTorch 的基礎數據類型是 torch.Tensor，該數據類型可以認為是 Pytorch 庫對 np.array 進行了更高一層的封裝。所以，在寫拓展程序時，其接口函數所需要的數據類型以及調用的庫會有些區別，下面會詳細解釋。

4.1. C++ Extensions For Python

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

import cv2import torch  # 不能刪掉, 因為需要動態加載torch的一些動態庫,后面會詳細說明.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一一對應.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 文件中調用了 affine_opencv 函數，而 affine_opencv 的 C++ 實現在 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 拓展，為了更加規范，這里在拓展接口程序（orbbec/warpaffine/src/warpaffine_ext.cpp）中通過 PYBIND11_MODULE 定義好接口，如下所示：

#include <torch/extension.h>#include<pybind11/numpy.h>
// python的C++拓展函數申明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++拓展函數申明(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拓展函數申明(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宏進一步封裝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 數組與 pybind11 的 py::array_t 相互對應，也即 Python 接口函數中，傳入的 np.array 數組，在 C++ 對應的函數中用 py::array_t 接收，操作 Numpy 數組，需要引入頭文件。

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

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

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

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

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

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

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

編譯和安裝成功

pip list 顯示相關安裝包信息