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++扩展的例子：

#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++中编写自定义的微分方程是一项具有挑战性的任务，并且需要领域专家知识。

#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文件中传递的名称定义。

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`。

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模型运行。

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 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)。