PyTorch C++ Extension on AMD GPU --- ROCm Blogs
本文演示了如何使用PyTorch C++扩展,并通过示例讨论了它相对于常规PyTorch模块的优势。实验在AMD GPU和ROCm 5.7.0软件上进行。有关支持的GPU和操作系统的更多信息,请参阅系统要求(Linux)。
介绍
由于易用性和模型的广泛可用性,PyTorch已成为机器学习从业者和爱好者的首选开发框架。PyTorch还允许您通过创建`torch.nn.Module`的派生类来轻松定制模型,这减少了与可微性相关的重复代码的需要。简而言之,PyTorch提供了广泛的支持。
但如果您想加速自定义模型呢?PyTorch提供了C++扩展来加速您的工作负载。这些扩展有优势:
• 它们为源外操作(PyTorch中不可用的操作)提供了一个快速的C++测试台,并且可以轻松集成到PyTorch模块中。
• 它们可以快速编译模型,无论是在CPU还是GPU上,只需一个附加的构建文件来编译C++模块。
PyTorch的自定义C++和CUDA扩展教程由Peter Goldsborough编写,这篇文章解释了PyTorch C++扩展如何减少模型的编译时间。PyTorch建立在一个C++后端之上,实现快速的计算操作。然而,构建PyTorch C++扩展的方式与PyTorch本身的构建方式不同。您可以在您的C++文件中包含PyTorch的库(torch.h
),以充分利用PyTorch的`tensor`和`Variable`接口,同时使用原生的C++库,如`iostream`。下面的代码片段是从PyTorch教程中取的使用C++扩展的例子:
cpp
#include <torch/extension.h>
#include <iostream>
torch::Tensor d_sigmoid(torch::Tensor z) {
auto s = torch::sigmoid(z);
return (1 - s) * s;
}
_d_sigmoid_函数计算了sigmoid函数的导数,并在后向传播中使用。您可以看到,实现是PyTorch的一个C++扩展。例如,`d_sigmoid`函数的返回值数据类型以及函数参数`z`是`torch::Tensor`。这是因为`torch/extension.h`头文件包含了著名的`ATen`张量计算库。让我们现在看看如何通过查看一个完整的示例来使用C++扩展来加速程序。
实现
在本节中,我们将在原生PyTorch和PyTorch C++中测试一个具有一个隐藏层的通用MLP网络。源代码受到了Peter的LLTM(长期记忆模型)示例的启发,我们为我们的MLP模型建立了类似的流程。
现在让我们在C++中实现_mlp_forward_和_mlp_backward_函数。PyTorch有`torch.autograd.Function`来在后台实现后向传递。PyTorch C++扩展要求我们在C++中定义后向传递,然后将它们绑定到PyTorch的`autograd`函数中。
如下所示,_mlp_forward_函数执行与MLP Python类中的计算相同,_mlp_backward_函数实现了输出相对于输入的导数。如果您对理解数学推导感兴趣,可以查看Prof. Tony Jebara的ML幻灯片中定义的_反向传播_部分中的后向传递方程。他代表了一个有两个隐藏层的MLP网络,并详细说明了后向传播的微分方程。为了简单起见,我们的示例中只考虑了一个隐藏层。请注意,在C++中编写自定义的微分方程是一项具有挑战性的任务,并且需要领域专家知识。
cpp
#include <torch/extension.h>
#include <vector>
#include <iostream>
torch::Tensor mlp_forward(
torch::Tensor input,
torch::Tensor hidden_weights,
torch::Tensor hidden_bias,
torch::Tensor output_weights,
torch::Tensor output_bias) {
// Compute the input/hidden layer
auto hidden = torch::addmm(hidden_bias, input, hidden_weights.t());
hidden = torch::relu(hidden);
// Compute the output layer
auto output = torch::addmm(output_bias, hidden, output_weights.t());
// Return the output
return output;
}
std::vector<torch::Tensor> mlp_backward(
torch::Tensor input,
torch::Tensor hidden_weights,
torch::Tensor hidden_bias,
torch::Tensor output_weights,
torch::Tensor output_bias,
torch::Tensor grad_output) {
// Compute the input/hidden layer
auto hidden = torch::addmm(hidden_bias, input, hidden_weights.t());
hidden = torch::relu(hidden);
// Compute the output layer
auto output = torch::addmm(output_bias, hidden, output_weights.t());
// Compute the gradients for output layer
auto grad_output_weights = torch::mm(grad_output.t(), hidden);
auto grad_output_bias = torch::sum(grad_output, /*dim=*/0).unsqueeze(0);
// Compute the gradients for input/hidden layer using chain rule
auto grad_hidden = torch::mm(grad_output, output_weights);
// grad_hidden = grad_hidden
auto grad_hidden_weights = torch::mm(grad_hidden.t(), input);
auto grad_hidden_bias = torch::sum(grad_hidden, /*dim=*/0).unsqueeze(0);
// Compute the gradients for input
auto grad_input = torch::mm(grad_hidden , hidden_weights);
// Return the gradients
return {grad_input, grad_hidden_weights, grad_hidden_bias, grad_output_weights, grad_output_bias};
}
以下是将C++实现使用`ATen`的Python绑定函数封装起来的示例。`PYBIND11_MODULE`将关键字_forward_映射到`mlp_forward`函数的指针,以及_backward_映射到`mlp_backward`函数。这将C++实现绑定到Python定义。宏`TORCH_EXTENSION_NAME`将在构建时在setup.py文件中传递的名称定义。
cpp
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("forward", &mlp_forward, "MLP forward");
m.def("backward", &mlp_backward, "MLP backward");
}
接下来,编写一个`setup.py`文件,导入`setuptools`库来帮助编译C++代码。要构建并安装C++扩展,运行`python setup.py install`命令。该命令会创建所有与`mlp.cpp`文件相关的构建文件,并提供一个可以导入到PyTorch模块中的模块`mlp_cpp`。
python
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CppExtension
setup(
name='mlp_cpp',
ext_modules=[
CppExtension('mlp_cpp', ['mlp.cpp']),
],
cmdclass={
'build_ext': BuildExtension
})
现在,让我们使用`torch.nn.Module`和`torch.autograd.Function`的帮助,准备一个由C++函数驱动的PyTorch的MLP类。这允许以更符合PyTorch原生方式使用C++函数。在下面的示例中,_MLP_类的forward函数指向`MLPFunction`的forward函数,它又指向C++的`mlp_forward`函数。这个信息流建立了一个工作流程,可以无缝地作为常规的PyTorch模型运行。
python
import math
from torch import nn
from torch.autograd import Function
import torch
import mlp_cpp
torch.manual_seed(42)
class MLPFunction(Function):
@staticmethod
def forward(ctx, input, hidden_weights, hidden_bias, output_weights, output_bias):
output = mlp_cpp.forward(input, hidden_weights, hidden_bias, output_weights, output_bias)
variables = [input, hidden_weights, hidden_bias, output_weights, output_bias]
ctx.save_for_backward(*variables)
return output
@staticmethod
def backward(ctx, grad_output):
grad_input, grad_hidden_weights, grad_hidden_bias, grad_output_weights, grad_output_bias = mlp_cpp.backward( *ctx.saved_variables, grad_output)
return grad_input, grad_hidden_weights, grad_hidden_bias, grad_output_weights, grad_output_bias
class MLP(nn.Module):
def __init__(self, input_features=5, hidden_features=15):
super(MLP, self).__init__()
self.input_features = input_features
self.hidden_weights = nn.Parameter(torch.rand(hidden_features,input_features))
self.hidden_bias = nn.Parameter(torch.rand(1, hidden_features))
self.output_weights = nn.Parameter(torch.rand(1,hidden_features))
self.output_bias = nn.Parameter(torch.rand(1, 1))
self.reset_parameters()
def reset_parameters(self):
stdv = 0.001
for weight in self.parameters():
weight.data.uniform_(-stdv, +stdv)
def forward(self, input):
return MLPFunction.apply(input, self.hidden_weights, self.hidden_bias, self.output_weights, self.output_bias)
现在,让我们使用 [trainer.py](PyTorch C++ Extension on AMD GPU --- ROCm Blogs) 来测试前向和后向计算的速度,并将原生 PyTorch 实现与 C++ 实现进行比较。
注意:在某些情况下,在进行基准测试以期望看到速度提升的趋势之前,你可能需要多次运行程序。
python
python trainer.py py
Forward: 0.102 milliseconds (ms) | Backward 0.223 milliseconds (ms)
我们可以看到,在 100,000 次运行中,原生 PyTorch 模型的平均前向传递时间是 0.102 毫秒,而对于 C++ 模型,它只需要 0.0904 毫秒(提升约 8%)。如果后向传递没有遵循相同的趋势,其实现可能没有优化。如前所述,将数学微分方程转换为 C++ 代码是一项具有挑战性的任务。随着模型的复杂性和大小的增加,在两个实验之间我们可能会看到更大的差异,正如 Peter 的 LLTM 示例中所注释的。尽管有一些实现挑战,C++ 正在证明它的速度更快,而且与 PyTorch 集成也更方便。
完整代码可以在 [src](PyTorch C++ Extension on AMD GPU --- ROCm Blogs) 文件夹中找到,它包含了以下结构:
-
[setup.py](https://rocm.blogs.amd.com/_downloads/1e638f7ade5de8f2cc73cd9f4ca07e54/setup.py) - 编译 C++ 模块的构建文件
-
[mlp.cpp](https://rocm.blogs.amd.com/_downloads/72080e8113297740e24fb96f8fe46b65/mlp.cpp) - C++ 模块
-
[mlp_cpp_train.py](https://rocm.blogs.amd.com/_downloads/00f3258c26bf3c8838dc72eb3a6ded8a/mlp_cpp_train.py) - 将 C++ 扩展应用于 PyTorch 模型
-
[mlp_train.py](https://rocm.blogs.amd.com/_downloads/65248a2373711bbdef8139c524f96a28/mlp_train.py) - 用于对比的原生 PyTorch 实现
-
[trainer.py](https://rocm.blogs.amd.com/_downloads/0d2415a09361672c52a5736a414ff5eb/trainer.py) - 用于测试 PyTorch 与 PyTorch 的 C++ 扩展的训练文件。
结论
这个博客通过一个使用自定义 PyTorch C++ 扩展的例子逐步向你演示。我们观察到,与原生 PyTorch 实现相比,自定义 C++ 扩展提高了模型的性能。这些扩展易于实现,并且可以轻松地插入到 PyTorch 模块中,预编译的开销很小。
此外,PyTorch 的 Aten
库为我们提供了大量功能,可以导入到 C++ 模块中,并模仿 PyTorch 风格的代码。总的来说,PyTorch C++ 扩展易于实现,是测试自定义操作在 CPU 和 GPU 上的性能的一个很好的选择。
致谢
我们想要感谢 Peter Goldsborough,因为他写了一篇非常精彩的[文章](Custom C++ and CUDA Extensions --- PyTorch Tutorials 2.3.0+cu121 documentation)。