Custom C++ and CUDA Extensions

0. Abstract

经历了一波 pybind11 和 CUDA 编程的学习, 接下来看一看 PyTorch 官方给的 C++/CUDA 扩展的教程. 发现极其简单, 就是直接用 setuptools 导出 PyTorch C++ 版代码的 Python 接口就可以了. 所以, 本博客包含以下内容:

LibTorch 初步;
C++ Extension 例子;

1. `LibTorch` 初步

在 PyTorch 的首页安装指引中就可以看到 PyTorch 是支持 C++/Java 的:

下载后解压到一个地方, 如 /opt/libtorch. 然后就可以使用 C++ 编写 PyTorch 程序了. 官方给的有相关例子, 我们选择最经典的 MNIST 手写数字识别项目来看一看:

mnist/
├── CMakeLists.txt
├── README.md
└── mnist.cpp

1.1 CMake 项目

CMakeLists.txt 是构建 cpp 项目的说明文件:

cmake_minimum_required(VERSION 3.5)
project(mnist)
set(CMAKE_CXX_STANDARD 17)find_package(Torch REQUIRED)option(DOWNLOAD_MNIST "Download the MNIST dataset from the internet" ON)
if (DOWNLOAD_MNIST)message(STATUS "Downloading MNIST dataset")execute_process(COMMAND python ${CMAKE_CURRENT_LIST_DIR}/../tools/download_mnist.py -d ${CMAKE_BINARY_DIR}/dataERROR_VARIABLE DOWNLOAD_ERROR)if (DOWNLOAD_ERROR)message(FATAL_ERROR "Error downloading MNIST dataset: ${DOWNLOAD_ERROR}")endif()
endif()add_executable(mnist mnist.cpp)
target_compile_features(mnist PUBLIC cxx_range_for)
target_link_libraries(mnist ${TORCH_LIBRARIES})

为了下载 MNIST 数据集, 这里用到了一个 Python 文件 ../tools/download_mnist.py, 执行 cmake 后, 编译根目录(build)会出现一个 data 数据文件夹.

find_package(Torch REQUIRED) 查找 libtorch 时可能需要指定路径:
find_package(Torch REQUIRED PATHS "path/to/libtorch/")
make 时, Ubuntu18.04 下出现错误: undefined reference to symbol ‘pthread_create@@GLIBC_2.2.5’.
=> 经查阅资料, 说: pthread 不是 linux 下的默认的库, 也就是在链接的时候, 无法找到 phread 库中线程函数的入口地址, 于是链接会失败.
=> 解决方案: target_link_libraries(mnist ${TORCH_LIBRARIES} -lpthread -lm)

make 之后, 执行 ./mnist 就能进行训练与测试了:

CUDA available! Training on GPU.
Train Epoch: 1 [59584/60000] Loss: 0.2078
Test set: Average loss: 0.2062 | Accuracy: 0.935
Train Epoch: 2 [59584/60000] Loss: 0.2039
Test set: Average loss: 0.1304 | Accuracy: 0.959
...

1.2 PyTorch C++ API

接下来看 C++ 代码:

struct Net : torch::nn::Module
{Net() : conv1(torch::nn::Conv2dOptions(1, 10, /*kernel_size=*/5)),conv2(torch::nn::Conv2dOptions(10, 20, /*kernel_size=*/5)),fc1(320, 50),fc2(50, 10){register_module("conv1", conv1);register_module("conv2", conv2);register_module("conv2_drop", conv2_drop);register_module("fc1", fc1);register_module("fc2", fc2);}torch::Tensor forward(torch::Tensor &x){x = torch::relu(torch::max_pool2d(conv1->forward(x), 2));x = torch::relu(torch::max_pool2d(conv2_drop->forward(conv2->forward(x)), 2));x = x.view({-1, 320});x = torch::relu(fc1->forward(x));x = torch::dropout(x, /*p=*/0.5, /*training=*/is_training());x = fc2->forward(x);return torch::log_softmax(x, /*dim=*/1);}torch::nn::Conv2d conv1;torch::nn::Conv2d conv2;torch::nn::Dropout2d conv2_drop;torch::nn::Linear fc1;torch::nn::Linear fc2;
};template<typename DataLoader>
void train(size_t epoch,Net &model,torch::Device device,DataLoader &data_loader,torch::optim::Optimizer &optimizer,size_t dataset_size
)
{model.train();size_t batch_idx = 0;for (auto &batch: data_loader){auto data = batch.data.to(device), targets = batch.target.to(device);auto output = model.forward(data);auto loss = torch::nll_loss(output, targets);AT_ASSERT(!std::isnan(loss.template item<float>()));optimizer.zero_grad();loss.backward();optimizer.step();...}
}

可以看到, 代码非常简单, 几乎和 Python 接口一致, 如果把 :: 换成 ., 就更像了. 不一样的是多了些类型限制以及一些语法. 具体的我们不多研究, 终究还是没有 Python 简洁好用. 但简单了解一下 PyTorch C++ API 的文档说明还是有必要的:

所以, 这个 LibTorch 既能用来写 C++ 项目, 也能用来给 PyTorch 写扩展. 不过官方还是推荐使用 Python 接口:

2. `C++ Extension` 例子

官方文档给的例子比较复杂, 这里举一个简单的例子, 把计算:

y = torch.relu(torch.matmul(x, w.t()) + b)

整合到一个操作里, 也就是使用 LibTorch C++ 编写一个等价的运算, 并导出 Python 接口. 这么做的理由是:

大概意思就是 Python 比较慢, 由 Python 一次次调用操作而频繁启动 CUDA 核会拖慢速度.

其实我觉得只有用 CUDA 编程把序列操作整合起来才能真正减少 CUDA 核的频繁启动, LibTorch 能加速可能就是因为 C++ 更快而已.

直接上代码吧, 整个项目的解构是这样子的:

LinearAct/
├── linearfun.py
├── linearact.cpp
└── setup.py

linearact.cpp 包含了组合操作的 forward 过程和 backward 过程, 前者计算正向的正常计算, 后者计算反向的梯度计算:

#include <torch/extension.h>  // 注意这里头文件和直接写 C++ 项目不一样
#include <vector>std::vector<at::Tensor> forward(torch::Tensor &input, torch::Tensor &weight, torch::Tensor &bias)
{auto relu_input = input.mm(weight.t()) + bias;auto output = torch::relu(relu_input);return {relu_input, output};  // relu_input 会在梯度计算时用到
}std::vector<torch::Tensor>
backward(torch::Tensor &grad_output, torch::Tensor &relu_input, torch::Tensor &input, torch::Tensor &weight)
{   // 求导链式法则auto grad_relu = grad_output.masked_fill(relu_input < 0, 0);auto grad_input = grad_relu.mm(weight);auto grad_weight = grad_relu.t().mm(input);auto grad_bias = grad_relu.sum(0);return {grad_input, grad_weight, grad_bias};
}PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {m.def("forward", &forward, "Custom forward");m.def("backward", &backward, "Custom backward");
}

这涉及到 pybind11 的用法, 详情见《pybind11 学习笔记》, 还涉及到使用 torch.autograd.Function 自定义运算的梯度计算, 详情见《PyTorch 中的 apply [autograd.Function]》. 总之, 现在我们使用 LibTorch 写了组合操作, 并写了其参数的梯度计算. linearfun.py 是利用 torch.autograd.Function 将 forward 和 backward 整合到一起, 组成一个完整的可以进行反向梯度传播的组合运算:

import torch  # 注意, 导入 linearact 前, 应先导入 torch
import linearactclass LinearActFunction(torch.autograd.Function):@staticmethoddef forward(ctx, input, weights, bias):relu_input, output = linearact.forward(input, weights, bias)  # c++ 函数variables = [relu_input, input, weights]ctx.save_for_backward(*variables)return output@staticmethoddef backward(ctx, grad_output):outputs = linearact.backward(grad_output, *ctx.saved_tensors)  # c++ 函数grad_x, grad_w, grad_b = outputsreturn grad_x, grad_w, grad_bmylinear = LinearActFunction.apply

LibTorch C++ 代码由 setuptools 导出 Python 接口:

from setuptools import setup
from torch.utils import cpp_extensionsetup(name='linearact',ext_modules=[cpp_extension.CppExtension('linearact', ['linearact.cpp'])],cmdclass={'build_ext': cpp_extension.BuildExtension}  # 整合了 pybind11 的功能
)

在命令行执行:

python setup.py install

就可以将 linearact 包安装到 Python 系统中, 任务完成. 下面进行验证:

import torch
from linearfun import mylinearx = torch.randn(2, 3, requires_grad=True)
w = torch.randn(2, 3, requires_grad=True)
b = torch.randn(2, requires_grad=True)
# 复制一份一样的参数
x1 = torch.from_numpy(x.detach().numpy())
w1 = torch.from_numpy(w.detach().numpy())
b1 = torch.from_numpy(b.detach().numpy())
x1.requires_grad_(True)
w1.requires_grad_(True)
b1.requires_grad_(True)# %% pytorch
y = torch.relu(torch.matmul(x, w.t()) + b)
y = y.norm(p=2)
print(y)y.backward()
print(x.grad)
print(w.grad)
print(b.grad)# %% custom
print('---------------------------')
y = mylinear(x1, w1, b1)
y = y.norm(p=2)
print(y)y.backward()
print(x1.grad)
print(w1.grad)
print(b1.grad)

执行一次:

tensor(1.2664, grad_fn=<LinalgVectorNormBackward0>)
tensor([[ 0.0851, -1.0418,  0.3958],[ 0.0566, -0.6925,  0.2631]])
tensor([[ 0.0000,  0.0000,  0.0000],[-1.0724,  0.3669, -0.1399]])
tensor([0.0000, 1.3864])
---------------------------
tensor(1.2664, grad_fn=<LinalgVectorNormBackward0>)
tensor([[ 0.0851, -1.0418,  0.3958],[ 0.0566, -0.6925,  0.2631]])
tensor([[ 0.0000,  0.0000,  0.0000],[-1.0724,  0.3669, -0.1399]])
tensor([0.0000, 1.3864])