备注
PyTorch Vulkan后端不再维护。请查看`ExecuTorch Vulkan Delegate <https://pytorch.org/executorch/stable/native-delegates-executorch-vulkan-delegate.html>`_ 的实现。
PyTorch Vulkan后端用户工作流程¶
Created On: Oct 26, 2020 | Last Updated: Oct 16, 2024 | Last Verified: Nov 05, 2024
作者: Ivan Kobzarev
简介¶
PyTorch 1.7支持在支持Vulkan图形和计算API的GPU上运行模型推理。主要目标设备是安卓设备上的移动GPU。Vulkan后端也可用于Linux、Mac和Windows桌面版本，以使用诸如英特尔集成GPU这样的Vulkan设备。该功能处于原型阶段，可能会发生变化。
构建带Vulkan后端的PyTorch¶
Vulkan后端默认不包括在内。包括Vulkan后端的主开关是cmake选项``USE_VULKAN``，可以通过环境变量``USE_VULKAN``设置。
要使用带Vulkan后端的PyTorch，需要从源码构建并设置额外配置。从GitHub主分支检出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 构建¶
使用 Vulkan 后端为指定的 ANDROID_ABI 构建 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 后端的算子。
默认情况下，optimize_for_mobile 的 backend='vulkan' 会重写图形，使得输入转到 Vulkan 后端，输出转到 CPU 后端，因此模型可以在 CPU 输入上运行并生成 CPU 输出。若关闭此功能，请在 optimize_for_mobile 中添加参数 optimization_blocklist={MobileOptimizerType.VULKAN_AUTOMATIC_GPU_TRANSFER}。(MobileOptimizerType 可以从 torch.utils.mobile_optimizer 导入)
查看更多信息，请参阅 torch.utils.mobile_optimizer 的 API 文档 <https://pytorch.org/docs/stable/mobile_optimizer.html>。
代码中使用 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¶
在加载模型时，我们需要指定模型将在 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。
有关使用 Vulkan 后端的示例可以在 PyTorch 仓库中的测试应用程序找到: https://github.com/pytorch/pytorch/blob/master/android/test_app/app/src/main/java/org/pytorch/testapp/MainActivity.java#L133
构建带 Vulkan 的安卓测试应用¶
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' 的应用程序。
无需上传到安卓设备即可测试模型¶
Vulkan 软件实现（例如 https://swiftshader.googlesource.com/SwiftShader ）可用于测试模型是否可以使用 PyTorch Vulkan 后端运行（例如检查是否支持所有模型运算符）。
