博客

在Ubuntu 20.04系统上完成

参考官方教程（坑比较多） https://pytorch.org/tutorials/prototype/vulkan_workflow.html

PyTorch的Vulkan backend坑非常多，不支持4维以上的Tensor（不大改网络的话三维卷积应该是没得玩了），缺少算子导致连MobileNet v2都跑不起来（这个让人很诧异，不知道是不是我编译的时候没用shaderc导致少了点东西），把ONNX模型迁移到Tensorflow Lite可能坑少一点？

1. 编译Vulkan SDK

下载Vulkan SDK并解压，参考官网的要求（

LunarXchange

；Getting Started – Tarball部分）安装各种依赖；在解压后的文件夹里执行./vulkansdk 进行编译，如果使用shaderc runtime的话执行./vulkansdk shaderc

2. 获取PyTorch源码以及其中的submodule，选择1.7以上的版本的release branch（例如release/1.10，直接从master branch的代码开始编译可能会有莫名其妙的错误）

git clone -b release/1.10 --recursive https://github.com/pytorch/pytorch
cd pytorch
git submodule sync
git submodule update --init --recursive

3. 编译PyTorch（参考

https://github.com/pytorch/pytorch#from-source

），在PyTorch源码目录中执行如下命令：

（编译之前需要将mkl和mkl-include更新到最新版本，不用shaderc的话将环境变量USE_VULKAN_SHADERC_RUNTIME的值设置为0）

conda install astunparse numpy ninja pyyaml mkl mkl-include setuptools cmake cffi typing_extensions future six requests dataclasses
# CUDA only: Add LAPACK support for the GPU if needed
conda install -c pytorch magma-cuda110  # or the magma-cuda* that matches your CUDA version from https://anaconda.org/pytorch/repo
USE_VULKAN=1 USE_VULKAN_SHADERC_RUNTIME=1 USE_VULKAN_WRAPPER=0 VULKAN_SDK=<vulkan sdk的路径> python setup.py install

4. 编译Android库

下载Android NDK，执行如下命令

ANDROID_NDK=<Android NDK的路径> ANDROID_ABI=arm64-v8a USE_VULKAN=1 sh ./scripts/build_android.sh

如果提示缺少vulkan_wrapper.h，需要将vulkan SDK中的部分C++头文件和源文件复制到NDK中

cp -v <Vulkan SDK的路径>/source/Vulkan-Tools/common/* <Android NDK的路径>/sources/third_party/vulkan/src/common/

在PyTorch源码根目录执行如下命令，编译示例Android工程：

USE_VULKAN=1 ANDROID_HOME=<Android SDK的路径> ANDROID_NDK=<Android NDK的路径> bash ./scripts/build_pytorch_android.sh <要编译的ABI类型>

其中，ABI类型可以包括多个ABI，例如 “armeabi-v7a;arm64-v8a”，默认的ABI似乎只有x86，至少需要把armeabi-v7a加上

5. 在Android工程中引入PyTorch的Java API以及预编译的JNI库：

引入示例Android工程中的Java文件（org.pytorch包）：

（<PyTorch源码目录>/android/pytorch_android/src/main/java）

将示例Android工程中编译好的静态链接库（<PyTorch源码目录>/android/pytorch_android/src/main/jniLibs）以及（<PyTorch源码目录>/android/pytorch_android/build/intermediates/cmake/release/obj）中的动态链接库复制到自己的Android工程中的app/src/main/jniLibs目录下

6. 在Android设备上加载模型并通过Vulkan在GPU上进行推理：

（将模型转换为lite模型可参考

Android | PyTorch

）

Module mModule = LiteModuleLoader.load("<pytorch lite模型的路径>", null, Device.VULKAN);
FloatBuffer inTensorBuffer = Tensor.allocateFloatBuffer(1 * 3 * 224 * 224);
long[] shape = {1, 3, 224, 224};        // 输入的tensor的shape
Tensor inputTensor = Tensor.fromBlob(inTensorBuffer, shape);
Tensor outputTensor = mModule.forward(IValue.from(inputTensor)).toTensor();