PyTorch C++扩展用于AMD GPU

24 篇文章 1 订阅
20 篇文章 0 订阅

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

  • 29
    点赞
  • 18
    收藏
    觉得还不错? 一键收藏
  • 打赏
    打赏
  • 5
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 5
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

109702008

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值