本文来自社区投稿
作者：教 主
原文链接：
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 显示相关安装包信息
*博客内容为网友个人发布，仅代表博主个人观点，如有侵权请联系工作人员删除。
