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}/data
ERROR_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 linearact
class LinearActFunction(torch.autograd.Function):
@staticmethod
def 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
@staticmethod
def backward(ctx, grad_output):
outputs = linearact.backward(grad_output, *ctx.saved_tensors) # c++ 函数
grad_x, grad_w, grad_b = outputs
return grad_x, grad_w, grad_b
mylinear = LinearActFunction.apply
LibTorch C++ 代码由 setuptools 导出 Python 接口:
from setuptools import setup
from torch.utils import cpp_extension
setup(
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 mylinear
x = 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])
可以看见两者一模一样. 至于测速什么的不在本博文的考虑范围之内, 只是想了解 PyTorch 如何进行 C++ 扩展.
