PyTorch C++扩展用于AMD GPU

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) 文件夹中找到,它包含了以下结构:

结论

这个博客通过一个使用自定义 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)。

相关推荐
小蜗子4 分钟前
Multi‐modal knowledge graph inference via media convergenceand logic rule
人工智能·知识图谱
SpikeKing17 分钟前
LLM - 使用 LLaMA-Factory 微调大模型 环境配置与训练推理 教程 (1)
人工智能·llm·大语言模型·llama·环境配置·llamafactory·训练框架
凌云行者26 分钟前
OpenGL入门005——使用Shader类管理着色器
c++·cmake·opengl
凌云行者29 分钟前
OpenGL入门006——着色器在纹理混合中的应用
c++·cmake·opengl
黄焖鸡能干四碗1 小时前
信息化运维方案,实施方案,开发方案,信息中心安全运维资料(软件资料word)
大数据·人工智能·软件需求·设计规范·规格说明书
1 小时前
开源竞争-数据驱动成长-11/05-大专生的思考
人工智能·笔记·学习·算法·机器学习
ctrey_1 小时前
2024-11-4 学习人工智能的Day21 openCV(3)
人工智能·opencv·学习
~yY…s<#>1 小时前
【刷题17】最小栈、栈的压入弹出、逆波兰表达式
c语言·数据结构·c++·算法·leetcode
攻城狮_Dream1 小时前
“探索未来医疗:生成式人工智能在医疗领域的革命性应用“
人工智能·设计·医疗·毕业
学习前端的小z2 小时前
【AIGC】如何通过ChatGPT轻松制作个性化GPTs应用
人工智能·chatgpt·aigc