在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();