深度解決添加復(fù)雜數(shù)據(jù)增強導(dǎo)致訓(xùn)練模型耗時長的痛點（2）
加入技術(shù)交流群
編譯完成后，可以運行 tools/collect_env.py，查看當(dāng)前一些必要工具的版本等一系列信息，輸出如下：
sys.platform : win32Python : 3.6.13 |Anaconda, Inc.| (default, Mar 16 2021, 11:37:27) [MSC v.1916 64 bit (AMD64)]CUDA available : TrueCUDA_HOME : C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.1NVCC : Not AvailableGPU 0 : NVIDIA GeForce GTX 1650OpenCV : 3.4.0PyTorch : 1.5.0PyTorch compiling details : PyTorch built with: - C++ Version: 199711 - MSVC 191627039 - Intel(R) Math Kernel Library Version 2020.0.0 Product Build 20191125 for Intel(R) 64 architecture applications - Intel(R) MKL-DNN v0.21.1 (Git Hash 7d2fd500bc78936d1d648ca713b901012f470dbc) - OpenMP 200203 - CPU capability usage: AVX2 - CUDA Runtime 10.1 - NVCC architecture flags: -gencode;arch=compute_37,code=sm_37;-gencode;arch=compute_50,code=sm_50;-gencode;arch=compute_60,code=sm_60;-gencode;arch=compute_61,code=sm_61;-gencode;arch=compute_70,code=sm_70;-gencode;arch=compute_75,code=sm_75;-gencode;arch=compute_37,code=compute_37 - CuDNN 7.6.4 - Magma 2.5.2 - Build settings: BLAS=MKL, BUILD_TYPE=Release, CXX_FLAGS=/DWIN32 /D_WINDOWS /GR /w /EHa /bigobj -openmp -DNDEBUG -DUSE_FBGEMM, PERF_WITH_AVX=1, PERF_WITH_AVX2=1, PERF_WITH_AVX512=1, USE_CUDA=ON, USE_EXCEPTION_PTR=1, USE_GFLAGS=OFF, USE_GLOG=OFF, USE_MKL=ON, USE_MKLDNN=ON, USE_MPI=OFF, USE_NCCL=OFF, USE_NNPACK=OFF, USE_OPENMP=ON, USE_STATIC_DISPATCH=OFF,
TorchVision : 0.6.0C/C++ Compiler : MSVC 191627045CUDA Compiler : 10.1
在運行 scripts/test_warpaffine_opencv.py 文件之前，由于 warpaffine_opencv.cpp 源碼用到相關(guān) opencv 庫，因此，還需要配置動態(tài)庫路徑，Windows 系統(tǒng)配置如下：
Windows 相關(guān)環(huán)境配置（opencv 第三方庫）
Linux 系統(tǒng)同樣也需要配置進行配置，命令如下：
root@aistation:/xxx/code/python_cpp_extension# export LD_LIBRARY_PATH=/xxx/code/python_cpp_extension/3rdparty/opencv/linux/libroot@aistation:/xxx/code/python_cpp_extension# ldconfig
也可以通過修改 ~/.bashrc 文件，加入上述 export LD_LIBRARY_PATH=/...，然后命令：source ~/.bashrc。也可以直接修改配置文件 /etc/profile，與修改 .bashrc 文件 一樣，對所有用戶有效。
可以通過 tools 下的 Dependencies_x64_Release 工具（運行：DependenciesGui.exe），查看編譯好的文件（.pyd）依賴的動態(tài)庫是否都配置完好，如下圖所示：
檢查編譯好的動態(tài)庫依賴的動態(tài)庫路徑
可以發(fā)現(xiàn)，該工具沒有找到 python36.dll、c10.dll、torch_cpu.dll、torch_python.dll 和 c10_cuda.dll 的路徑。
這里說明一下，Python 相關(guān)的 dll 庫以及 torch 相關(guān)的動態(tài)庫是動態(tài)加載的，也就是說，如果你在 Python 代碼中寫一句：import torch，只有在程序運行時才會動態(tài)加載 torch 相關(guān)庫。
所以，Dependencies_x64_Release工具檢查不到編譯好的 warpaffine_ext.cp36-win_amd64.pyd 文件依賴完好性。
這里還需要說明一下為什么 warpaffine_ext.cp36-win_amd64.pyd 需要依賴 torch 相關(guān)庫，這是因為源文件 orbbec/warpaffine/src/warpaffine_ext.cpp 兼容了 PyTorch 的 C++ 拓展，所以依賴 torch 和 cuda 相關(guān)動態(tài)庫文件，如果你單純只在 orbbec/warpaffine/src/warpaffine_ext.cpp 實現(xiàn)純粹 Python 的 C++拓展，則是不需要依賴 torch 和 cuda 相關(guān)動態(tài)庫。
配置好之后，還需要將 warpaffine_ext.cp36-win_amd64.pyd 無法動態(tài)加載的動態(tài)庫文件（opencv_world453.dll）放到 scripts/test_warpaffine_opencv.py 同路徑之下（Linux 系統(tǒng)也一樣)，如下圖所示：
拷貝動態(tài)庫與測試腳本同一目錄
需要注意一個問題，有時候，如果在 docker 中進行編譯和安裝，其最終生成的 Python 安裝包（.egg）文件并不會安裝到當(dāng)前 Python 環(huán)境下的 site-packages 中。
也就意味著，在 Python 文件中執(zhí)行：from orbbec.warpaffine import affine_opencv 會失敗。
原因是 orbbec.warpaffine 并不在其 Python 的搜索路徑中，這個時候有兩種解決辦法：一種是在執(zhí)行：python setup.py install 時，加上 --prefix='install path'，但是經(jīng)過本人驗證，有時候不可行，另外一種辦法是在 Python 文件中，將 orbbec 文件夾路徑添加到 Python 的搜索路徑中，如下所示：
import cv2import torch # 不能刪掉, 因為需要動態(tài)加載torch的一些動態(tài)庫.import numpy as np
# 添加下述兩行代碼，這里默認此python腳本所在目錄的上一層目錄路徑包含orbbec文件夾._FILE_PATH = os.path.dirname(os.path.abspath(__file__))sys.path.insert(0, os.path.join(_FILE_PATH, "../"))
from orbbec.warpaffine import affine_opencv # C++ interface
4.2. C++/CUDA Extensions For PyTorch
PyTorch 的 C++/CUDA 拓展同樣也是利用 Pybind11 工具，但是，由于 PyTorch 使用的基礎(chǔ)數(shù)據(jù)類型是 torch.Tensor 類型，因此，在寫拓展程序中，必須要有 libtorch 庫中對應(yīng)的數(shù)據(jù)類型與 PyTorch 的 tensor 類型對應(yīng)，這樣才能進行正確傳參。這里需要知道 PyTorch 對應(yīng)的 C++ 版本 ibtorch 中幾個常用的庫和命名空間。
常用的命名空間：
at(ATen) 負責(zé)聲明和定義 Tensor 運算，是最常用到的命名空間；
c10 是 ATen 的基礎(chǔ)，包含了 PyTorch 的核心抽象、Tensor 和 Storage 數(shù)據(jù)結(jié)構(gòu)的實際實現(xiàn)
torch 命名空間下定義的 Tensor 相比于 ATen 增加自動求導(dǎo)功能
PyTorch 的 Aten 目錄下的主要構(gòu)成：
ATen（ATen 核心源文件）
TH（Torch 張量計算庫）
THC（Torch CUDA 張量計算庫）
THCUNN（Torch CUDA 神經(jīng)網(wǎng)絡(luò)庫）
THNN（Torch 神經(jīng)網(wǎng)絡(luò)庫）
C10 是 Caffe Tensor Library 的縮寫。這里存放的都是最基礎(chǔ)的 Tensor 庫的代碼，可以運行在服務(wù)端和移動端，C10 主要目的之一是為了統(tǒng)一 PyTorch 的張量計算后端代碼和 caffe2 的張量計算后端代碼。
libtorch 中還有個 csrc 模塊，主要適用于 C++ 和 Python 的 API 之間的相互映射，比如 PyTorch 的 nn.Conv2d 對應(yīng)于 torch 中的 at:conv2d，其次是 autograd 和自動求導(dǎo)機制。
了解如上內(nèi)容后，首先來看 Python 測試代碼，如下所示（scripts/test_warpaffine_torch_cpu.py）：
import cv2import torchimport numpy as npfrom orbbec.warpaffine import affine_torch # C++ interface
data_path = "demo.png"
img = cv2.imread(data_path)# transform img(numpy.array) to tensor(torch.Tensor)# use permuteimg_tensor = torch.from_numpy(img / 255.0).permute(2, 0, 1).contiguous()img_tensor = img_tensor.unsqueeze(0).float()
src_tensor = torch.tensor([[38.29, 51.69, 1.0], [73.53, 51.69, 1.0], [56.02, 71.73, 1.0]], dtype=torch.float32).unsqueeze(0)dst_tensor = torch.tensor([[262.0, 324.0], [325.0, 323.0], [295.0, 349.0]], dtype=torch.float32).unsqueeze(0)
# compute affine transform matrixmatrix_l = torch.transpose(src_tensor, 1, 2).bmm(src_tensor)matrix_l = torch.inverse(matrix_l)matrix_r = torch.transpose(src_tensor, 1, 2).bmm(dst_tensor)affine_matrix = torch.transpose(matrix_l.bmm(matrix_r), 1, 2)
warpffine_img = affine_torch(img_tensor, affine_matrix, 112, 112)
warpffine_img = warpffine_img.squeeze(0).permute(1, 2, 0).numpy()cv2.imwrite("torch_affine_cpu.png", np.uint8(warpffine_img * 255.0))
從上述代碼可以看到，Python 文件中調(diào)用了 affine_torch 函數(shù)，并且傳入的參數(shù)類型是 cpu 類型的 tensor，而 affine_torch 的 C++ 實現(xiàn)在 orbbec/warpaffine/src/warpaffine_ext.cpp 中，如下所示：
#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宏進一步封裝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");}
從上述代碼可以看出，根據(jù)宏 WITH_CUDA 和 tensor 類型控制 affine_torch 最終底層執(zhí)行 affine_cpu 還是 affine_gpu 函數(shù)。同時也注意到，Python 中的 torch.Tensor 類型與 libtorch 庫中的 at::Tensor 對應(yīng)。再看看 affine_cpu 函數(shù)的具體實現(xiàn)（orbbec/warpaffine/src/cpu/warpaffine_torch_v2.cpp）：
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){ at::Tensor result; // AT_DISPATCH_FLOATING_TYPES: input.scalar_type() => scalar_t AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "affine_cpu", [&] { result = affine_torch_cpu<scalar_t>(input, affine_matrix, out_h, out_w); }); return result;}進一步看 affine_torch_cpu 函數(shù)的具體實現(xiàn)（orbbec/warpaffine/src/cpu/warpaffine_torch_v2.cpp）：
template <typename scalar_t>at::Tensor affine_torch_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) { AT_ASSERTM(input.device().is_cpu(), "input must be a CPU tensor"); AT_ASSERTM(affine_matrix.device().is_cpu(), "affine_matrix must be a CPU tensor");
auto matrix_ptr = affine_matrix.contiguous().data_ptr<scalar_t>(); auto input_ptr = input.contiguous().data_ptr<scalar_t>(); auto nimgs = input.size(0); auto img_c = input.size(1); auto img_h = input.size(2); auto img_w = input.size(3); auto in_img_size = img_c * img_h * img_w; auto out_img_size = img_c * out_h * out_w;
// build dst tensor auto output_tensor = at::zeros({nimgs, img_c, out_h, out_w}, input.options()); auto output_ptr = output_tensor.contiguous().data_ptr<scalar_t>(); for(int i = 0; i < nimgs; i++) { scalar_t* matrix = matrix_ptr + i * 6; scalar_t* in = input_ptr + i * in_img_size; scalar_t* out = output_ptr + i * out_img_size; affine_cpu_kernel<scalar_t>(img_h, img_w, img_c, img_w*img_h, out_h, out_w, out_h*out_w, out, in, matrix, 0.0f); }
return output_tensor;}
這里有一個非常注意的地方就是，上述代碼中的 tensor 的 .contiguous() 方法（上述代碼第 10、11、21 行）。
可以看到，我們在獲取 tensor 的數(shù)據(jù)指針時候（data_ptr()），PyTorch 官方示例代碼和 MMDtection/MMCV 中的一些相關(guān)代碼都推薦先做這個操作。
這是因為，不管是在 Python 還是在 C++ 代碼中，使用 permute()、transpose()、view() 等方法操作返回一個新的 tensor 時，其與舊的 tensor 是共享數(shù)據(jù)存儲，所以他們的 storage 不會發(fā)生變化，只是會重新返回一個新的 view，這樣做的目的是減少數(shù)據(jù)拷貝，減少內(nèi)存消耗，一定程度上加速網(wǎng)絡(luò)訓(xùn)練或推理過程，如果在 Python 端對 tensor 做了 .contiguous() 操作，則在 C++ 端就不需要再做了，因為 .contiguous() 是一個深拷貝操作。
permute 操作分析
接下來，再來看 PyTorch 的 CUDA 擴展，首先測試文件 test_warpaffine_torch_gpu.py 如下：
import cv2import torchimport numpy as npfrom orbbec.warpaffine import affine_torch # CUDA interface
data_path = "demo.png"
img = cv2.imread(data_path)# transform img(numpy.array) to tensor(torch.Tensor)# use permuteimg_tensor = torch.from_numpy(img / 255.0).permute(2, 0, 1).contiguous()img_tensor = img_tensor.unsqueeze(0).float()img_tensor = img_tensor.cuda() # gpu tensor
# dst -> srcsrc_tensor = torch.tensor([[38.29, 51.69, 1.0], [73.53, 51.69, 1.0], [56.02, 71.73, 1.0]], dtype=torch.float32).unsqueeze(0)dst_tensor = torch.tensor([[262.0, 324.0], [325.0, 323.0], [295.0, 349.0]], dtype=torch.float32).unsqueeze(0)src_tensor = src_tensor.cuda() # gpu tensordst_tensor = dst_tensor.cuda() # gpu tensor
# compute affine transform matrixmatrix_l = torch.transpose(src_tensor, 1, 2).bmm(src_tensor)matrix_l = torch.inverse(matrix_l)matrix_r = torch.transpose(src_tensor, 1, 2).bmm(dst_tensor)affine_matrix = torch.transpose(matrix_l.bmm(matrix_r), 1, 2)affine_matrix = affine_matrix.contiguous().cuda() # gpu tensor
warpffine_img = affine_torch(img_tensor, affine_matrix, 112, 112)warpffine_img = warpffine_img.cpu().squeeze(0).permute(1, 2, 0).numpy()cv2.imwrite("torch_affine_gpu.png", np.uint8(warpffine_img * 255.0))
從上述腳本代碼可以看到，affine_torch 接收的是 GPU 類型的Tensor 數(shù)據(jù)，其底層會在 GPU 上執(zhí)行相關(guān)計算。進一步分析 orbbec/warpaffine/src/warpaffine_ext.cpp 中的 affine_torch() 函數(shù)的 CUDA 接口，可以發(fā)現(xiàn)，最終調(diào)用的是 affine_gpu() 函數(shù)，如下代碼所示：
at::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){ CHECK_INPUT(input); CHECK_INPUT(affine_matrix);
// Ensure CUDA uses the input tensor device. at::DeviceGuard guard(input.device());
return affine_cuda_forward(input, affine_matrix, out_h, out_w);}
可以發(fā)現(xiàn)，最終執(zhí)行的是 affine_cuda_forward() 函數(shù)，如下代碼所示：
at::Tensor affine_cuda_forward(const at::Tensor& input, /*[B, C, H, W]*/ const at::Tensor& affine_matrix, /*[B, 2, 3]*/ const int out_h, const int out_w){ // build dst tensor auto nimgs = input.size(0); auto img_c = input.size(1); auto img_h = input.size(2); auto img_w = input.size(3); const int output_size = nimgs * img_c * out_h * out_w; auto output_tensor = at::zeros({nimgs, img_c, out_h, out_w}, input.options());
AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "affine_cuda", [&] { auto matrix_ptr = affine_matrix.data_ptr<scalar_t>(); auto input_ptr = input.data_ptr<scalar_t>(); auto output_ptr = output_tensor.data_ptr<scalar_t>();
// launch kernel function on GPU with CUDA. affine_gpu_kernel<scalar_t><<<GET_BLOCKS(output_size), THREADS_PER_BLOCK, 0, at::cuda::getCurrentCUDAStream()>>>(output_size, img_h, img_w, img_c, out_h, out_w, output_ptr, input_ptr, matrix_ptr, 0.0f); });
return output_tensor;}
通過配置 grid_size 和 block_size 之后，啟動核函數(shù): affine_gpu_kernel，關(guān)于核函數(shù)這一部分涉及很多 CUDA 知識，這里并不進行展開說明。最終返回 GPU 類型的 output_tensor 給 Python 接口。
5. GPU-Accelerated Augmentation
在掌握了 PyTorch 的 C++/CUDA 拓展之后，我們就可以輕松做到與 NVIDIA 的 DALI 庫一樣的加速效果，不管多么復(fù)雜的數(shù)據(jù)增強，都可以通過上述操作進行一定程度上的加速，偽代碼如下所示（假設(shè)編譯和安裝步驟都已完成）：
for _, (img, local_labels) in enumerate(train_loader): global_step += 1 # 這里假設(shè)從train_loader取出的gpu類型的Tensor, 如果是cpu類型的Tensor, 則需要首先放到對應(yīng)的編號為：local_rank的GPU上. # local_rank = torch.distributed.get_rank() # ================== add data augmentation (這里只做一個示意)=================== batch = img.shape[0] # get batchsize devive = img.device # get local_rank src_tensor = torch.tensor([[38.29, 51.69, 1.0], [73.53, 51.69, 1.0], [56.02, 71.73, 1.0]],dtype=torch.float32).unsqueeze(0) dst_tensor = torch.tensor([[42.0, 52.0], [78.0, 55.0], [58.0, 74.0]], dtype=torch.float32).unsqueeze(0) src_tensor = src_tensor.repeat(batch, 1, 1) dst_tensor = dst_tensor.repeat(batch, 1, 1) # compute affine transform matrix matrix_l = torch.transpose(src_tensor, 1, 2).bmm(src_tensor) matrix_l = torch.inverse(matrix_l) matrix_r = torch.transpose(src_tensor, 1, 2).bmm(dst_tensor) affine_matrix = torch.transpose(matrix_l.bmm(matrix_r), 1, 2) affine_matrix = affine_matrix.contiguous().to(devive) # python端做了.contiguous()操作, 則CUDA拓展底層不需要再做. img = affine_torch(img, affine_matrix, 112, 112) # 在gpu上進行數(shù)據(jù)增強 # ============================================================================== local_embeddings = backbone(img) loss: torch.Tensor = module_partial_fc(local_embeddings, local_labels, opt)
【注】：本工程編譯完成后，可以將 orbbec 文件夾直接拷貝自己的訓(xùn)練工程，然后在對應(yīng)的需要調(diào)用拓展函數(shù)的 Python 文件中（比如上述代碼在 train.py 文件中），通過之前提到的方法，將 orbbec 文件夾所在路徑加入到 Python 環(huán)境中，就可以正常調(diào)用拓展函數(shù)了（比如：affine_torch）。
Reference：
[1]: Custom C++ and CUDA Extensions
[2]: https://github.com/NVIDIA/DALI
[3]:https://github.com/open-mmlab/mmdetection/tree/v2.0.0
[4]: GitHub - open-mmlab/mmcv: OpenMMLab Computer Vision Foundation
[5]: GitHub - openppl-public/ppl.cv: ppl.cv is a high-performance image processing library of openPPL supporting various platforms.
[6]: https://github.com/pytorch/extension-cpp
[7]: Keywords - setuptools 65.6.0.post20221119 documentation
[8]: mmdetection源碼剖析(1)--NMS
[9]: JeffWang：教程：Python中使用C++/CUDA｜以PointNet中的ball query 為例
[10]: OpenMMLab：PyTorch 源碼解讀之 cpp_extension：揭秘 C++/CUDA 算子實現(xiàn)和調(diào)用全流程
[11]: Pytorch拓展進階(二)：Pytorch結(jié)合C++以及Cuda拓展 - Oldpan的個人博客
[12]: https://docs.python.org/zh-cn/3/extending/building.html
[13]: 王炳明：花了兩天，終于把 Python 的 setup.py 給整明白了
[14]: 【pybind11】--python C/C++擴展編譯
[15]: pizh12thu：Python/C++混合編程利器Pybind11實踐
*博客內(nèi)容為網(wǎng)友個人發(fā)布，僅代表博主個人觀點，如有侵權(quán)請聯(lián)系工作人員刪除。
