4.2. C++/CUDA Extensions For PyTorch
PyTorch 的 C++/CUDA 拓展同样也是利用 Pybind11 工具，但是，由于 PyTorch 使用的基础数据类型是 torch.Tensor 类型，因此，在写拓展程序中，必须要有 libtorch 库中对应的数据类型与 PyTorch 的 tensor 类型对应，这样才能进行正确传参。这里需要知道 PyTorch 对应的 C++ 版本 ibtorch 中几个常用的库和命名空间。
常用的命名空间：
at(ATen) 负责声明和定义 Tensor 运算，是最常用到的命名空间；
c10 是 ATen 的基础，包含了 PyTorch 的核心抽象、Tensor 和 Storage 数据结构的实际实现
torch 命名空间下定义的 Tensor 相比于 ATen 增加自动求导功能
PyTorch 的 Aten 目录下的主要构成：
ATen（ATen 核心源文件）
TH（Torch 张量计算库）
THC（Torch CUDA 张量计算库）
THCUNN（Torch CUDA 神经网络库）
THNN（Torch 神经网络库）
C10 是 Caffe Tensor Library 的缩写。这里存放的都是最基础的 Tensor 库的代码，可以运行在服务端和移动端，C10 主要目的之一是为了统一 PyTorch 的张量计算后端代码和 caffe2 的张量计算后端代码。
libtorch 中还有个 csrc 模块，主要适用于 C++ 和 Python 的 API 之间的相互映射，比如 PyTorch 的 nn.Conv2d 对应于 torch 中的 at:conv2d，其次是 autograd 和自动求导机制。
了解如上内容后，首先来看 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 文件中调用了 affine_torch 函数，并且传入的参数类型是 cpu 类型的 tensor，而 affine_torch 的 C++ 实现在 orbbec/warpaffine/src/warpaffine_ext.cpp 中，如下所示：
#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");}
从上述代码可以看出，根据宏 WITH_CUDA 和 tensor 类型控制 affine_torch 最终底层执行 affine_cpu 还是 affine_gpu 函数。同时也注意到，Python 中的 torch.Tensor 类型与 libtorch 库中的 at::Tensor 对应。再看看 affine_cpu 函数的具体实现（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 函数的具体实现（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 的数据指针时候（data_ptr()），PyTorch 官方示例代码和 MMDtection/MMCV 中的一些相关代码都推荐先做这个操作。
这是因为，不管是在 Python 还是在 C++ 代码中，使用 permute()、transpose()、view() 等方法操作返回一个新的 tensor 时，其与旧的 tensor 是共享数据存储，所以他们的 storage 不会发生变化，只是会重新返回一个新的 view，这样做的目的是减少数据拷贝，减少内存消耗，一定程度上加速网络训练或推理过程，如果在 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 数据，其底层会在 GPU 上执行相关计算。进一步分析 orbbec/warpaffine/src/warpaffine_ext.cpp 中的 affine_torch() 函数的 CUDA 接口，可以发现，最终调用的是 affine_gpu() 函数，如下代码所示：
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);}
可以发现，最终执行的是 affine_cuda_forward() 函数，如下代码所示：
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 之后，启动核函数: affine_gpu_kernel，关于核函数这一部分涉及很多 CUDA 知识，这里并不进行展开说明。最终返回 GPU 类型的 output_tensor 给 Python 接口。
*博客内容为网友个人发布，仅代表博主个人观点，如有侵权请联系工作人员删除。
