注意
PyTorch Vulkan 后端不再维护。请查看 ExecuTorch Vulkan Delegate 实现。
PyTorch Vulkan 后端用户工作流程¶
创建日期：2020年10月26日 | 最后更新：2024年10月16日 | 最后验证：2024年11月5日
作者: 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。根据您从主分支检出代码的时间，您可能会看到不同的数字，但它应该大于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，获取默认的预训练浮点模型。
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 输出。要禁用此功能，请将参数 optimization_blocklist={MobileOptimizerType.VULKAN_AUTOMATIC_GPU_TRANSFER} 添加到 optimize_for_mobile 中。（MobileOptimizerType 可以从 torch.utils.mobile_optimizer 导入）
欲了解更多信息，请参阅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。
.vulkan() 函数在张量上调用时，会将张量复制到Vulkan设备上，并且对于使用此张量作为输入调用的操作符 - 操作符将在Vulkan设备上运行，其输出也将位于Vulkan设备上。
.cpu() 在Vulkan张量上调用的函数将把其数据复制到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。
使用Vulkan后端的示例可以在PyTorch仓库中的测试应用程序中找到： 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’的应用程序。
无需上传到安卓设备即可测试模型¶
Vulkan的软件实现（例如https://swiftshader.googlesource.com/SwiftShader）可用于测试模型是否可以使用PyTorch Vulkan后端运行（例如检查是否支持所有模型操作符）。
