快捷方式

ExecuTorch Vulkan 委托

ExecuTorch Vulkan 委托是基于跨平台 Vulkan GPU API 标准构建的 ExecuTorch 的原生 GPU 委托。它主要设计用于利用 GPU 加速 Android 设备上的模型推理,但可在支持 Vulkan 实现的任何平台上使用:笔记本电脑、服务器和边缘设备。

注意

Vulkan 委托目前正在积极开发中,其组件可能会发生变化。

什么是 Vulkan?

Vulkan 是作为 OpenGL 的继任者开发的低级 GPU API 规范。与之前的规范相比,它旨在为开发人员提供对 GPU 的更明确控制,以减少开销并最大化现代图形硬件的功能。

Vulkan 已被 GPU 供应商广泛采用,市场上大多数现代 GPU(包括台式机和移动设备)都支持 Vulkan。Vulkan 也从 Android 7.0 开始包含在 Android 中。

请注意,Vulkan 是 GPU API,而不是 GPU 数学库。也就是说,它提供了一种在 GPU 上执行计算和图形操作的方法,但没有内置的高性能计算内核库。

Vulkan 计算库

ExecuTorch Vulkan 委托是对称为Vulkan 计算库的独立运行时的包装。Vulkan 计算库的目标是通过 GLSL 计算着色器为 PyTorch 运算符提供 GPU 实现。

Vulkan 计算库是PyTorch Vulkan 后端的一个分支/迭代。PyTorch Vulkan 后端的核心组件已分叉到 ExecuTorch 中,并针对 AOT 图形模式的模型推理风格进行了调整(与采用急切执行模型推理风格的 PyTorch 相反)。

Vulkan 计算库的组件包含在 executorch/backends/vulkan/runtime/ 目录中。核心组件如下列出并进行描述

runtime/
├── api/ .................... Wrapper API around Vulkan to manage Vulkan objects
└── graph/ .................. ComputeGraph class which implements graph mode inference
    └── ops/ ................ Base directory for operator implementations
        ├── glsl/ ........... GLSL compute shaders
        │   ├── *.glsl
        │   └── conv2d.glsl
        └── impl/ ........... C++ code to dispatch GPU compute shaders
            ├── *.cpp
            └── Conv2d.cpp

功能

Vulkan 委托目前支持以下功能

  • 内存规划

    • 生命周期不重叠的中间张量将共享内存分配。这减少了模型推理的峰值内存使用量。

  • 基于功能的分区:

    • 可以通过分区器将图部分降低到 Vulkan 委托,分区器将识别 Vulkan 委托支持的节点(即运算符),并仅降低受支持的子图

  • 支持上限动态形状:

    • 只要张量的当前形状小于降低期间指定的边界,张量就可以在推理之间改变形状

除了增加运算符覆盖范围之外,以下功能目前正在开发中

  • 量化支持

    • 我们目前正在开发对 8 位动态量化的支持,并计划在未来扩展到其他量化方案。

  • 内存布局管理

    • 内存布局是优化性能的重要因素。我们计划引入图传递,在整个图中引入内存布局转换,以优化对内存布局敏感的运算符,例如卷积和矩阵乘法。

  • 选择性构建

    • 我们计划通过选择要构建的运算符/着色器来控制构建大小

端到端示例

要进一步了解 Vulkan 委托的功能以及如何使用它,请考虑以下使用 MobileNet V2 的端到端示例。

将模型编译并降低到 Vulkan 委托

假设已设置并安装 ExecuTorch,则可以使用以下脚本生成降低的 MobileNet V2 模型,作为 vulkan_mobilenetv2.pte

import torch
import torchvision.models as models

from torch.export import export, ExportedProgram
from torchvision.models.mobilenetv2 import MobileNet_V2_Weights
from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
from executorch.exir import EdgeProgramManager, ExecutorchProgramManager, to_edge
from executorch.exir.backend.backend_api import to_backend

mobilenet_v2 = models.mobilenetv2.mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT).eval()
sample_inputs = (torch.randn(1, 3, 224, 224), )

exported_program: ExportedProgram = export(mobilenet_v2, sample_inputs)
edge: EdgeProgramManager = to_edge(exported_program)

# Lower the model to Vulkan backend
edge = edge.to_backend(VulkanPartitioner())

exec_prog = edge.to_executorch()

with open("vulkan_mobilenetv2.pte", "wb") as file:
    exec_prog.write_to_file(file)

与其他 ExecuTorch 委托一样,可以使用 to_backend() API 将模型降低到 Vulkan 委托。Vulkan 委托实现了 VulkanPartitioner 类,该类识别图中 Vulkan 委托支持的节点(即运算符),并将模型的兼容部分分离出来在 GPU 上执行。

这意味着即使模型包含一些不受支持的运算符,也可以将其降低到 Vulkan 委托。这只是意味着图的仅一部分将在 GPU 上执行。

注意

可以检查Vulkan 分区器代码以检查 Vulkan 委托中当前实现了哪些运算符。

构建 Vulkan 委托库

构建和测试 Vulkan Delegate 的最简单方法是针对 Android 构建并在本地 Android 设备上进行测试。Android 设备内置对 Vulkan 的支持,并且 Android NDK 附带一个 GLSL 编译器,这是编译 Vulkan Compute Library 的 GLSL 计算着色器所必需的。

在使用 CMake 构建时,可以通过设置 -DEXECUTORCH_BUILD_VULKAN=ON 来构建 Vulkan Delegate 库。

首先,确保已安装 Android NDK - 建议使用 Android NDK r25c。还应安装 Android SDK,以便可以访问 adb

# Recommended version is Android NDK r25c.
export ANDROID_NDK=<path_to_ndk>
# Select an appropriate Android ABI
export ANDROID_ABI=arm64-v8a
# All subsequent commands should be performed from ExecuTorch repo root
cd <path_to_executorch_root>
# Make sure adb works
adb --version

使用 Vulkan Delegate 构建和安装 ExecuTorch 库(适用于 Android)

# From executorch root directory
(rm -rf cmake-android-out && \
  pp cmake . -DCMAKE_INSTALL_PREFIX=cmake-android-out \
    -DCMAKE_TOOLCHAIN_FILE=$ANDROID_NDK/build/cmake/android.toolchain.cmake \
    -DANDROID_ABI=$ANDROID_ABI \
    -DEXECUTORCH_BUILD_VULKAN=ON \
    -DPYTHON_EXECUTABLE=python \
    -Bcmake-android-out && \
  cmake --build cmake-android-out -j16 --target install)

在设备上运行 Vulkan 模型

注意

由于目前运算符支持受限,只有二进制算术运算符将在 GPU 上运行。由于大多数运算符都是通过可移植运算符执行的,因此推理速度较慢。

现在,部分委托模型可以在您设备的 GPU 上(部分)执行了!

# Build a model runner binary linked with the Vulkan delegate libs
cmake --build cmake-android-out --target vulkan_executor_runner -j32

# Push model to device
adb push vulkan_mobilenetv2.pte /data/local/tmp/vulkan_mobilenetv2.pte
# Push binary to device
adb push cmake-android-out/backends/vulkan/vulkan_executor_runner /data/local/tmp/runner_bin

# Run the model
adb shell /data/local/tmp/runner_bin --model_path /data/local/tmp/vulkan_mobilenetv2.pte

文档

访问 PyTorch 的综合开发者文档

查看文档

教程

获取针对初学者和高级开发者的深入教程

查看教程

资源

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

查看资源