4.2. C++/CUDA Extensions For PyTorch

PyTorch 的 C++/CUDA 拓展同樣也是利用 Pybind11 工具，但是，由于 PyTorch 使用的基礎數(shù)據(jù)類型是 torch.Tensor 類型，因此，在寫拓展程序中，必須要有 libtorch 庫中對應的數(shù)據(jù)類型與 PyTorch 的 tensor 類型對應，這樣才能進行正確傳參。這里需要知道 PyTorch 對應的 C++ 版本 ibtorch 中幾個常用的庫和命名空間。

常用的命名空間：

• at(ATen) 負責聲明和定義 Tensor 運算，是最常用到的命名空間；

• c10 是 ATen 的基礎，包含了 PyTorch 的核心抽象、Tensor 和 Storage 數(shù)據(jù)結(jié)構(gòu)的實際實現(xiàn)

• torch 命名空間下定義的 Tensor 相比于 ATen 增加自動求導功能

PyTorch 的 Aten 目錄下的主要構(gòu)成：

• ATen（ATen 核心源文件）

• TH（Torch 張量計算庫）

• THC（Torch CUDA 張量計算庫）

• THCUNN（Torch CUDA 神經(jīng)網(wǎng)絡庫）

• THNN（Torch 神經(jīng)網(wǎng)絡庫）

C10 是 Caffe Tensor Library 的縮寫。這里存放的都是最基礎的 Tensor 庫的代碼，可以運行在服務端和移動端，C10 主要目的之一是為了統(tǒng)一 PyTorch 的張量計算后端代碼和 caffe2 的張量計算后端代碼。

libtorch 中還有個 csrc 模塊，主要適用于 C++ 和 Python 的 API 之間的相互映射，比如 PyTorch 的 nn.Conv2d 對應于 torch 中的 at:conv2d，其次是 autograd 和自動求導機制。

了解如上內(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 對應。再看看 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)絡訓練或推理過程，如果在 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 接口。