注意
PyTorch Vulkan 後端已停止維護。請改為檢閱 ExecuTorch Vulkan Delegate 實作。
PyTorch Vulkan 後端使用者工作流程¶
建立於:2020 年 10 月 26 日 | 最後更新:2024 年 10 月 16 日 | 最後驗證:2024 年 11 月 05 日
簡介¶
PyTorch 1.7 支援在支援 Vulkan 圖形和計算 API 的 GPU 上執行模型推論的能力。主要目標裝置是 Android 裝置上的行動 GPU。Vulkan 後端也可以用於 Linux、Mac 和 Windows 桌面版本,以使用 Vulkan 裝置,例如 Intel 整合型 GPU。此功能處於原型階段,可能會有所變更。
使用 Vulkan 後端建置 PyTorch¶
預設情況下不包含 Vulkan 後端。包含 Vulkan 後端的主要開關是 cmake 選項 USE_VULKAN
,可以透過環境變數 USE_VULKAN
設定。
若要使用具有 Vulkan 後端的 PyTorch,我們需要從原始碼使用其他設定來建置它。從 GitHub master 分支簽出 PyTorch 原始碼。
Vulkan 封裝器的選用用法¶
預設情況下,將使用 vulkan_wrapper 程式庫在執行時載入 Vulkan 程式庫。如果您指定環境變數 USE_VULKAN_WRAPPER=0
,則 libvulkan 將會直接連結。
桌面版本¶
Vulkan SDK¶
從 https://vulkan.lunarg.com/sdk/home 下載 VulkanSDK 並設定環境變數 VULKAN_SDK
將 VulkanSDK 解壓縮到 VULKAN_SDK_ROOT
資料夾,按照 VulkanSDK 指示為您的系統安裝 VulkanSDK。
適用於 Mac
cd $VULKAN_SDK_ROOT
source setup-env.sh
sudo python install_vulkan.py
建置 PyTorch
適用於 Linux
cd PYTORCH_ROOT
USE_VULKAN=1 USE_VULKAN_SHADERC_RUNTIME=1 USE_VULKAN_WRAPPER=0 python setup.py install
適用於 Mac
cd PYTORCH_ROOT
USE_VULKAN=1 USE_VULKAN_SHADERC_RUNTIME=1 USE_VULKAN_WRAPPER=0 MACOSX_DEPLOYMENT_TARGET=10.9 CC=clang CXX=clang++ python setup.py install
成功建置後,開啟另一個終端機並驗證已安裝的 PyTorch 版本。
import torch
print(torch.__version__)
在撰寫此文件時,版本為 1.8.0a0+41237a4。根據您從 master 簽出程式碼的時間,您可能會看到不同的數字,但它應該大於 1.7.0。
Android 版本¶
若要為指定的 ANDROID_ABI
建置具有 Vulkan 後端的 LibTorch for android。
cd PYTORCH_ROOT
ANDROID_ABI=arm64-v8a USE_VULKAN=1 sh ./scripts/build_android.sh
若要準備可以直接在您的應用程式中使用的 pytorch_android aars
cd $PYTORCH_ROOT
USE_VULKAN=1 sh ./scripts/build_pytorch_android.sh
模型準備¶
安裝 torchvision,取得預設的預先訓練 float 模型。
pip install torchvision
Python 腳本,將預先訓練的 mobilenet_v2 儲存到檔案
import torch
import torchvision
model = torchvision.models.mobilenet_v2(pretrained=True)
model.eval()
script_model = torch.jit.script(model)
torch.jit.save(script_model, "mobilenet2.pt")
PyTorch 1.7 Vulkan 後端僅支援 float 32 位元運算子。預設模型需要額外的步驟,以最佳化運算子融合
from torch.utils.mobile_optimizer import optimize_for_mobile
script_model_vulkan = optimize_for_mobile(script_model, backend='vulkan')
torch.jit.save(script_model_vulkan, "mobilenet2-vulkan.pt")
結果模型只能在 Vulkan 後端上使用,因為它包含特定於 Vulkan 後端的運算子。
預設情況下,optimize_for_mobile
與 backend='vulkan'
會重寫圖形,以便將輸入傳輸到 Vulkan 後端,並將輸出傳輸到 CPU 後端,因此,該模型可以在 CPU 輸入上執行並產生 CPU 輸出。若要停用此功能,請將引數 optimization_blocklist={MobileOptimizerType.VULKAN_AUTOMATIC_GPU_TRANSFER}
新增至 optimize_for_mobile
。(可以從 torch.utils.mobile_optimizer
匯入 MobileOptimizerType
)
如需更多資訊,請參閱 torch.utils.mobile_optimizer API 文件。
在程式碼中使用 Vulkan 後端¶
C++ API¶
at::is_vulkan_available()
auto tensor = at::rand({1, 2, 2, 3}, at::device(at::kCPU).dtype(at::kFloat));
auto tensor_vulkan = t.vulkan();
auto module = torch::jit::load("$PATH");
auto tensor_output_vulkan = module.forward(inputs).toTensor();
auto tensor_output = tensor_output.cpu();
at::is_vulkan_available()
函數嘗試初始化 Vulkan 後端,如果成功找到 Vulkan 裝置並建立內容,它將傳回 true,否則傳回 false。
在 Tensor 上呼叫的 .vulkan()
函數會將張量複製到 Vulkan 裝置,並且對於使用此張量作為輸入呼叫的運算子,該運算子將在 Vulkan 裝置上執行,並且其輸出將在 Vulkan 裝置上。
在 Vulkan 張量上呼叫的 .cpu()
函數會將其資料複製到 CPU 張量 (預設)
使用 Vulkan 裝置上的張量作為輸入呼叫的運算子將會在 Vulkan 裝置上執行。如果 Vulkan 後端不支援運算子,則會擲回例外。
支援的運算子清單
_adaptive_avg_pool2d
_cat
add.Scalar
add.Tensor
add_.Tensor
addmm
avg_pool2d
clamp
convolution
empty.memory_format
empty_strided
hardtanh_
max_pool2d
mean.dim
mm
mul.Scalar
relu_
reshape
select.int
slice.Tensor
transpose.int
transpose_
unsqueeze
upsample_nearest2d
view
這些運算子允許在 Vulkan 後端上使用 torchvision 模型進行影像分類。
Python API¶
torch.is_vulkan_available()
已公開到 Python API。
tensor.to(device='vulkan')
的運作方式與 .vulkan()
將張量移動到 Vulkan 裝置相同。
在撰寫本教學時,.vulkan()
尚未公開到 Python API,但計畫將其納入。
Android Java API¶
為了讓 Android API 在 Vulkan 後端上執行模型,我們必須在模型載入期間指定此項。
import org.pytorch.Device;
Module module = Module.load("$PATH", Device.VULKAN)
FloatBuffer buffer = Tensor.allocateFloatBuffer(1 * 3 * 224 * 224);
Tensor inputTensor = Tensor.fromBlob(buffer, new int[]{1, 3, 224, 224});
Tensor outputTensor = mModule.forward(IValue.from(inputTensor)).toTensor();
在這種情況下,所有輸入將會透明地從 CPU 複製到 Vulkan 裝置,並且模型將會在 Vulkan 裝置上執行,輸出也會透明地複製到 CPU。
可以在 PyTorch 儲存庫中的測試應用程式中找到使用 Vulkan 後端的範例:https://github.com/pytorch/pytorch/blob/master/android/test_app/app/src/main/java/org/pytorch/testapp/MainActivity.java#L133
使用 Vulkan 構建 Android 測試應用程式¶
cd $PYTORCH_ROOT
USE_VULKAN=1 sh ./scripts/build_pytorch_android.sh
或者,如果您只需要特定的 ABI,您可以將其設定為參數。
cd $PYTORCH_ROOT
USE_VULKAN=1 sh ./scripts/build_pytorch_android.sh $ANDROID_ABI
將準備好的模型 mobilenet2-vulkan.pt
新增到測試應用程式的 assets 目錄中。
cp mobilenet2-vulkan.pt $PYTORCH_ROOT/android/test_app/app/src/main/assets/
cd $PYTORCH_ROOT
gradle -p android test_app:installMbvulkanLocalBaseDebug
成功安裝後,即可在裝置上啟動名為“MBQ”的應用程式。
無需上傳到 Android 裝置即可測試模型¶
Vulkan 的軟體實作(例如 https://swiftshader.googlesource.com/SwiftShader)可用於測試模型是否可以使用 PyTorch Vulkan 後端執行(例如,檢查是否支援所有模型運算子)。