• 教程 >
  • PyTorch Vulkan 后端用户工作流程
快捷方式

注意

PyTorch Vulkan 后端不再维护。请改用 ExecuTorch Vulkan 委托 实现。

PyTorch Vulkan 后端用户工作流程

**作者**:Ivan Kobzarev

简介

PyTorch 1.7 支持在支持 Vulkan 图形和计算 API 的 GPU 上运行模型推理的功能。主要目标设备是 Android 设备上的移动 GPU。Vulkan 后端也可用于 Linux、Mac 和 Windows 桌面版本,以使用 Intel 集成 GPU 等 Vulkan 设备。此功能处于原型阶段,可能会发生变化。

使用 Vulkan 后端构建 PyTorch

默认情况下不包含 Vulkan 后端。包含 Vulkan 后端的主要开关是 cmake 选项 USE_VULKAN,可以通过环境变量 USE_VULKAN 设置。

要使用 Vulkan 后端的 PyTorch,我们需要使用其他设置从源代码构建它。从 GitHub master 分支检出 PyTorch 源代码。

Vulkan 包装器的可选用法

默认情况下,Vulkan 库将使用 vulkan_wrapper 库在运行时加载。如果指定环境变量 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 后端的 Android 版 LibTorch。

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,获取默认的预训练浮点模型。

pip install torchvision

将预训练的 mobilenet_v2 保存到文件的 Python 脚本

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 后端仅支持浮点 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 后端上使用。

默认情况下,带有 backend='vulkan'optimize_for_mobile 会重写图形,以便将输入传输到 Vulkan 后端,并将输出传输到 CPU 后端,因此,模型可以在 CPU 输入上运行并生成 CPU 输出。要禁用此功能,请将参数 optimization_blocklist={MobileOptimizerType.VULKAN_AUTOMATIC_GPU_TRANSFER} 添加到 optimize_for_mobile 中。(MobileOptimizerType 可以从 torch.utils.mobile_optimizer 中导入)

有关更多信息,请参阅 torch.utils.mobile_optimizerAPI 文档

在代码中使用 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。

在张量上调用 .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 添加到测试应用程序资源。

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 后端运行(例如,检查是否支持所有模型运算符)。

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取面向初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源