5. GPU-Accelerated Augmentation

在掌握了 PyTorch 的 C++/CUDA 拓展之后，我们就可以轻松做到与 NVIDIA 的 DALI 库一样的加速效果，不管多么复杂的数据增强，都可以通过上述操作进行一定程度上的加速，伪代码如下所示（假设编译和安装步骤都已完成）：

for _, (img, local_labels) in enumerate(train_loader):    global_step += 1    # 这里假设从train_loader取出的gpu类型的Tensor, 如果是cpu类型的Tensor, 则需要首先放到对应的编号为：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上进行数据增强    # ==============================================================================            local_embeddings = backbone(img)    loss: torch.Tensor = module_partial_fc(local_embeddings, local_labels, opt)

【注】：本工程编译完成后，可以将 orbbec 文件夹直接拷贝自己的训练工程，然后在对应的需要调用拓展函数的 Python 文件中（比如上述代码在 train.py 文件中），通过之前提到的方法，将 orbbec 文件夹所在路径加入到 Python 环境中，就可以正常调用拓展函数了（比如：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 算子实现和调用全流程

• [11]: Pytorch拓展进阶(二)：Pytorch结合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实践

本文仅做学术分享，如有侵权，请联系删文。

专栏文章内容及配图由作者撰写发布，仅供工程师学习之用，如有侵权或者其他违规问题，请联系本站处理。 联系我们