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

5. GPU-Accelerated Augmentation

在掌握了 PyTorch 的 C++/CUDA 拓展之后，我們就可以輕松做到與 NVIDIA 的 DALI 庫一樣的加速效果，不管多么復(fù)雜的數(shù)據(jù)增強(qiáng)，都可以通過上述操作進(jìn)行一定程度上的加速，偽代碼如下所示（假設(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上進(jìn)行數(shù)據(jù)增強(qiáng)    # ==============================================================================            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拓展進(jìn)階(二)：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++擴(kuò)展編譯

• [15]: pizh12thu：Python/C++混合編程利器Pybind11實踐

本文僅做學(xué)術(shù)分享，如有侵權(quán)，請聯(lián)系刪文。