陕西做网站公司,开平网站制作,至少保存十个以上域名网站,老网站绑定新网站如何做PyTorch C Extension on AMD GPU — ROCm Blogs 本文演示了如何使用PyTorch C扩展#xff0c;并通过示例讨论了它相对于常规PyTorch模块的优势。实验在AMD GPU和ROCm 5.7.0软件上进行。有关支持的GPU和操作系统的更多信息#xff0c;请参阅系统要求#xff08;Linux#xf…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 iostreamtorch::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 iostreamtorch::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::vectortorch::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 layerauto 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 layerauto 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 ruleauto grad_hidden torch::mm(grad_output, output_weights);// grad_hidden grad_hiddenauto 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 inputauto 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, CppExtensionsetup(namemlp_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 torchimport mlp_cpptorch.manual_seed(42)class MLPFunction(Function):staticmethoddef 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 outputstaticmethoddef 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_biasclass MLP(nn.Module):def __init__(self, input_features5, hidden_features15):super(MLP, self).__init__()self.input_features input_featuresself.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.001for 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 pyForward: 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.0cu121 documentation)。
