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实践
本文仅做学术分享，如有侵权，请联系删文。
*博客内容为网友个人发布，仅代表博主个人观点，如有侵权请联系工作人员删除。
发布文章
