PyTorch Vulkan Backend User Workflow ==================================== **Author**: `Ivan Kobzarev `_ Introduction ------------ PyTorch 1.7 supports the ability to run model inference on GPUs that support the Vulkan graphics and compute API. The primary target devices are mobile GPUs on Android devices. The Vulkan backend can also be used on Linux, Mac, and Windows desktop builds to use Vulkan devices like Intel integrated GPUs. This feature is in the prototype stage and is subject to change. Building PyTorch with Vulkan backend ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Vulkan backend is not included by default. The main switch to include Vulkan backend is cmake option ``USE_VULKAN``, that can be set by environment variable ``USE_VULKAN``. To use PyTorch with Vulkan backend, we need to build it from source with additional settings. Checkout the PyTorch source code from GitHub master branch. Optional usage of vulkan wrapper ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ By default, Vulkan library will be loaded at runtime using the vulkan_wrapper library. If you specify the environment variable ``USE_VULKAN_WRAPPER=0`` libvulkan will be linked directly. Desktop build ^^^^^^^^^^^^^ Vulkan SDK ^^^^^^^^^^ Download VulkanSDK from https://vulkan.lunarg.com/sdk/home and set environment variable ``VULKAN_SDK`` Unpack VulkanSDK to ``VULKAN_SDK_ROOT`` folder, install VulkanSDK following VulkanSDK instructions for your system. For Mac: :: cd $VULKAN_SDK_ROOT source setup-env.sh sudo python install_vulkan.py Building PyTorch: For Linux: :: cd PYTORCH_ROOT USE_VULKAN=1 USE_VULKAN_SHADERC_RUNTIME=1 USE_VULKAN_WRAPPER=0 python setup.py install For 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 After successful build, open another terminal and verify the version of installed PyTorch. :: import torch print(torch.__version__) At the time of writing of this recipe, the version is 1.8.0a0+41237a4. You might be seeing different numbers depending on when you check out the code from master, but it should be greater than 1.7.0. Android build ^^^^^^^^^^^^^ To build LibTorch for android with Vulkan backend for specified ``ANDROID_ABI``. :: cd PYTORCH_ROOT ANDROID_ABI=arm64-v8a USE_VULKAN=1 sh ./scripts/build_android.sh To prepare pytorch_android aars that you can use directly in your app: :: cd $PYTORCH_ROOT USE_VULKAN=1 sh ./scripts/build_pytorch_android.sh Model preparation ----------------- Install torchvision, get the default pretrained float model. :: pip install torchvision Python script to save pretrained mobilenet_v2 to a file: :: 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 backend supports only float 32bit operators. The default model needs additional step that will optimize operators fusing :: 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") The result model can be used only on Vulkan backend as it contains specific to the Vulkan backend operators. By default, ``optimize_for_mobile`` with ``backend='vulkan'`` rewrites the graph so that inputs are transferred to the Vulkan backend, and outputs are transferred to the CPU backend, therefore, the model can be run on CPU inputs and produce CPU outputs. To disable this, add the argument ``optimization_blocklist={MobileOptimizerType.VULKAN_AUTOMATIC_GPU_TRANSFER}`` to ``optimize_for_mobile``. (``MobileOptimizerType`` can be imported from ``torch.utils.mobile_optimizer``) For more information, see the `torch.utils.mobile_optimizer` `API documentation `_. Using Vulkan backend in code ---------------------------- 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()`` function tries to initialize Vulkan backend and if Vulkan device is successfully found and context is created - it will return true, false otherwise. ``.vulkan()`` function called on Tensor will copy tensor to Vulkan device, and for operators called with this tensor as input - the operator will run on Vulkan device, and its output will be on the Vulkan device. ``.cpu()`` function called on Vulkan tensor will copy its data to CPU tensor (default) Operators called with a tensor on a Vulkan device as an input will be executed on a Vulkan device. If an operator is not supported for the Vulkan backend the exception will be thrown. List of supported operators: :: _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 Those operators allow to use torchvision models for image classification on Vulkan backend. Python API ---------- ``torch.is_vulkan_available()`` is exposed to Python API. ``tensor.to(device='vulkan')`` works as ``.vulkan()`` moving tensor to the Vulkan device. ``.vulkan()`` at the moment of writing of this tutorial is not exposed to Python API, but it is planned to be there. Android Java API ---------------- For Android API to run model on Vulkan backend we have to specify this during model loading: :: 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(); In this case, all inputs will be transparently copied from CPU to the Vulkan device, and model will be run on Vulkan device, the output will be copied transparently to CPU. The example of using Vulkan backend can be found in test application within the PyTorch repository: https://github.com/pytorch/pytorch/blob/master/android/test_app/app/src/main/java/org/pytorch/testapp/MainActivity.java#L133 Building android test app with Vulkan ------------------------------------- 1. Build pytorch android with Vulkan backend for all android ABIs ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ :: cd $PYTORCH_ROOT USE_VULKAN=1 sh ./scripts/build_pytorch_android.sh Or if you need only specific abi you can set it as an argument: :: cd $PYTORCH_ROOT USE_VULKAN=1 sh ./scripts/build_pytorch_android.sh $ANDROID_ABI 2. Add vulkan model to test application assets ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Add prepared model ``mobilenet2-vulkan.pt`` to test applocation assets: :: cp mobilenet2-vulkan.pt $PYTORCH_ROOT/android/test_app/app/src/main/assets/ 3. Build and Install test applocation to connected android device ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ :: cd $PYTORCH_ROOT gradle -p android test_app:installMbvulkanLocalBaseDebug After successful installation, the application with the name 'MBQ' can be launched on the device. Testing models without uploading to android device -------------------------------------------------- Software implementations of Vulkan (e.g. https://swiftshader.googlesource.com/SwiftShader ) can be used to test if a model can be run using PyTorch Vulkan Backend (e.g. check if all model operators are supported).