做網(wǎng)站發(fā)票windows優(yōu)化大師好用嗎
PyTorch C++ Extension on AMD GPU — ROCm Blogs
本文演示了如何使用PyTorch C++擴展,并通過示例討論了它相對于常規(guī)PyTorch模塊的優(yōu)勢。實驗在AMD GPU和ROCm 5.7.0軟件上進行。有關(guān)支持的GPU和操作系統(tǒng)的更多信息,請參閱系統(tǒng)要求(Linux)。
介紹
由于易用性和模型的廣泛可用性,PyTorch已成為機器學(xué)習(xí)從業(yè)者和愛好者的首選開發(fā)框架。PyTorch還允許您通過創(chuàng)建`torch.nn.Module`的派生類來輕松定制模型,這減少了與可微性相關(guān)的重復(fù)代碼的需要。簡而言之,PyTorch提供了廣泛的支持。
但如果您想加速自定義模型呢?PyTorch提供了C++擴展來加速您的工作負(fù)載。這些擴展有優(yōu)勢:
? 它們?yōu)樵赐獠僮?#xff08;PyTorch中不可用的操作)提供了一個快速的C++測試臺,并且可以輕松集成到PyTorch模塊中。
? 它們可以快速編譯模型,無論是在CPU還是GPU上,只需一個附加的構(gòu)建文件來編譯C++模塊。
PyTorch的自定義C++和CUDA擴展教程由Peter Goldsborough編寫,這篇文章解釋了PyTorch C++擴展如何減少模型的編譯時間。PyTorch建立在一個C++后端之上,實現(xiàn)快速的計算操作。然而,構(gòu)建PyTorch C++擴展的方式與PyTorch本身的構(gòu)建方式不同。您可以在您的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_函數(shù)計算了sigmoid函數(shù)的導(dǎo)數(shù),并在后向傳播中使用。您可以看到,實現(xiàn)是PyTorch的一個C++擴展。例如,`d_sigmoid`函數(shù)的返回值數(shù)據(jù)類型以及函數(shù)參數(shù)`z`是`torch::Tensor`。這是因為`torch/extension.h`頭文件包含了著名的`ATen`張量計算庫。讓我們現(xiàn)在看看如何通過查看一個完整的示例來使用C++擴展來加速程序。
實現(xiàn)
在本節(jié)中,我們將在原生PyTorch和PyTorch C++中測試一個具有一個隱藏層的通用MLP網(wǎng)絡(luò)。源代碼受到了Peter的LLTM(長期記憶模型)示例的啟發(fā),我們?yōu)槲覀兊腗LP模型建立了類似的流程。
現(xiàn)在讓我們在C++中實現(xiàn)_mlp_forward_和_mlp_backward_函數(shù)。PyTorch有`torch.autograd.Function`來在后臺實現(xiàn)后向傳遞。PyTorch C++擴展要求我們在C++中定義后向傳遞,然后將它們綁定到PyTorch的`autograd`函數(shù)中。
如下所示,_mlp_forward_函數(shù)執(zhí)行與MLP?Python類中的計算相同,_mlp_backward_函數(shù)實現(xiàn)了輸出相對于輸入的導(dǎo)數(shù)。如果您對理解數(shù)學(xué)推導(dǎo)感興趣,可以查看Prof. Tony Jebara的ML幻燈片中定義的_反向傳播_部分中的后向傳遞方程。他代表了一個有兩個隱藏層的MLP網(wǎng)絡(luò),并詳細(xì)說明了后向傳播的微分方程。為了簡單起見,我們的示例中只考慮了一個隱藏層。請注意,在C++中編寫自定義的微分方程是一項具有挑戰(zhàn)性的任務(wù),并且需要領(lǐng)域?qū)<抑R。
#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 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++實現(xiàn)使用`ATen`的Python綁定函數(shù)封裝起來的示例。`PYBIND11_MODULE`將關(guān)鍵字_forward_映射到`mlp_forward`函數(shù)的指針,以及_backward_映射到`mlp_backward`函數(shù)。這將C++實現(xiàn)綁定到Python定義。宏`TORCH_EXTENSION_NAME`將在構(gòu)建時在setup.py文件中傳遞的名稱定義。
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {m.def("forward", &mlp_forward, "MLP forward");m.def("backward", &mlp_backward, "MLP backward");
}
接下來,編寫一個`setup.py`文件,導(dǎo)入`setuptools`庫來幫助編譯C++代碼。要構(gòu)建并安裝C++擴展,運行`python setup.py install`命令。該命令會創(chuàng)建所有與`mlp.cpp`文件相關(guān)的構(gòu)建文件,并提供一個可以導(dǎo)入到PyTorch模塊中的模塊`mlp_cpp`。
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CppExtensionsetup(name='mlp_cpp',ext_modules=[CppExtension('mlp_cpp', ['mlp.cpp']),],cmdclass={'build_ext': BuildExtension})
現(xiàn)在,讓我們使用`torch.nn.Module`和`torch.autograd.Function`的幫助,準(zhǔn)備一個由C++函數(shù)驅(qū)動的PyTorch的MLP類。這允許以更符合PyTorch原生方式使用C++函數(shù)。在下面的示例中,_MLP_類的forward函數(shù)指向`MLPFunction`的forward函數(shù),它又指向C++的`mlp_forward`函數(shù)。這個信息流建立了一個工作流程,可以無縫地作為常規(guī)的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 output@staticmethoddef 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_features=5, hidden_features=15):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)
現(xiàn)在,讓我們使用 [trainer.py](PyTorch C++ Extension on AMD GPU — ROCm Blogs) 來測試前向和后向計算的速度,并將原生 PyTorch 實現(xiàn)與 C++ 實現(xiàn)進行比較。
注意:在某些情況下,在進行基準(zhǔn)測試以期望看到速度提升的趨勢之前,你可能需要多次運行程序。
python trainer.py pyForward: 0.102 milliseconds (ms) | Backward 0.223 milliseconds (ms)
我們可以看到,在 100,000 次運行中,原生 PyTorch 模型的平均前向傳遞時間是 0.102 毫秒,而對于 C++ 模型,它只需要 0.0904 毫秒(提升約 8%)。如果后向傳遞沒有遵循相同的趨勢,其實現(xiàn)可能沒有優(yōu)化。如前所述,將數(shù)學(xué)微分方程轉(zhuǎn)換為 C++ 代碼是一項具有挑戰(zhàn)性的任務(wù)。隨著模型的復(fù)雜性和大小的增加,在兩個實驗之間我們可能會看到更大的差異,正如 Peter 的 LLTM 示例中所注釋的。盡管有一些實現(xiàn)挑戰(zhàn),C++ 正在證明它的速度更快,而且與 PyTorch 集成也更方便。
完整代碼可以在 [src](PyTorch C++ Extension on AMD GPU — ROCm Blogs) 文件夾中找到,它包含了以下結(jié)構(gòu):
- [setup.py](https://rocm.blogs.amd.com/_downloads/1e638f7ade5de8f2cc73cd9f4ca07e54/setup.py) - 編譯 C++ 模塊的構(gòu)建文件
- [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++ 擴展應(yīng)用于 PyTorch 模型
- [mlp_train.py](https://rocm.blogs.amd.com/_downloads/65248a2373711bbdef8139c524f96a28/mlp_train.py) - 用于對比的原生 PyTorch 實現(xiàn)
- [trainer.py](https://rocm.blogs.amd.com/_downloads/0d2415a09361672c52a5736a414ff5eb/trainer.py) - 用于測試 PyTorch 與 PyTorch 的 C++ 擴展的訓(xùn)練文件。
結(jié)論
這個博客通過一個使用自定義 PyTorch C++ 擴展的例子逐步向你演示。我們觀察到,與原生 PyTorch 實現(xiàn)相比,自定義 C++ 擴展提高了模型的性能。這些擴展易于實現(xiàn),并且可以輕松地插入到 PyTorch 模塊中,預(yù)編譯的開銷很小。
此外,PyTorch 的?Aten
?庫為我們提供了大量功能,可以導(dǎo)入到 C++ 模塊中,并模仿 PyTorch 風(fēng)格的代碼??偟膩碚f,PyTorch C++ 擴展易于實現(xiàn),是測試自定義操作在 CPU 和 GPU 上的性能的一個很好的選擇。
致謝
我們想要感謝 Peter Goldsborough,因為他寫了一篇非常精彩的[文章](Custom C++ and CUDA Extensions — PyTorch Tutorials 2.3.0+cu121 documentation)。