[mobile] Mobile Perf Recipe

Author: IvanKobzarev

recipes_source/recipes/mobile_perf.py

@@ -0,0 +1,188 @@
+"""
+Pytorch Mobile Performance Guideline 
+====================================
+
+Introduction
+------------
+Performance (aka latency) is crucial to most, if not all,
+applications and use-cases of ML model inference on mobile devices.
+Today, PyTorch executes the models on the CPU backend pending availability
+of other hardware backends such as GPU, DSP, and NPU.
+
+"""
+
+######################################################################
+# Model preparation
+# -----------------
+#
+# Next recipes you can take (offline) while preparing the model
+# to have an optimized model that will probably have shorter execution time
+# (higher performance, lower latency) on the mobile device.
+
+######################################################################
+# 1. Use torch.utils.mobile_optimizer
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+#
+# Torch mobile_optimizer package does several optimizations with the model,
+# which will help to conv2d and linear operations.
+# It pre-packs model weights in an optimized format and fuses ops above with relu
+# if it is the next operation.
+
+from torch.utils.mobile_optimizer import optimize_for_mobile
+traced_model = torch.jit.load("input_model_path")
+optimized_model = optimize_for_mobile(traced_model)
+torch.jit.save(optimized_model, "output_model_path")
+
+######################################################################
+# 2. Fuse operators using torch.quantization.fuse_modules
+# Do not be confused that fuse_modules is in the quantization package.
+# It works for all types of torch script modules.
+# `torch.quantization.fuse_modules<https://pytorch.org/docs/stable/quantization.html#torch.quantization.fuse_modules>`
+# fuses a list of modules into a single module.
+# It fuses only the following sequence of modules:
+#
+# - Convolution, Batch normalization
+# - Convolution, Batch normalization, Relu
+# - Convolution, Relu
+# - Linear, Relu
+# 
+# This script shows how to fuse Batch Normalization in torchvision [mobileNetV2](https://github.com/pytorch/vision/blob/master/torchvision/models/mobilenet.py#L33).
+#
+# ConvBNReLU: nn.Sequential[nn.Conv2d, nn.BatchNorm2d, nn.ReLU]
+#
+# InvertedResidual: nn.Sequential[ConvBnReLU, ConvBNReLU, nn.Conv2d, nn.BatchNorm2d]
+#
+
+import torch
+import torchvision
+
+def fuse_model(model):
+    for idx, m in enumerate(model.modules()):
+        name = m._get_name();
+        if name == 'ConvBNReLU':
+            torch.quantization.fuse_modules(m, ['0', '1'], inplace=True)
+        if name == 'InvertedResidual':
+            for idx in range(len(m.conv)):
+                if type(m.conv[idx]) == torch.nn.Conv2d:
+                    torch.quantization.fuse_modules(m.conv, [str(idx), str(idx + 1)], inplace=True)
+
+m = torchvision.models.mobilenet_v2(pretrained=True)
+m.eval()
+fuse_model(m)
+torch.jit.trace(m, torch.rand(1, 3, 224, 224)).save("mobilenetV2-bnfused.pt")
+
+######################################################################
+#
+# Quantization
+# ------------
+#
+
+######################################################################
+# 3. Try a quantized version of your model
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+#
+# You can find more about PyTorch quantization in `the dedicated tutorial<https://pytorch.org/blog/introduction-to-quantization-on-pytorch/>`.
+#
+# Quantization of the model not only moves computation to int8, 
+# but also reduces the size of your model on a disk.
+# That size reduction helps to reduce disk read operations during the first load of the model and decreases the amount of RAM.
+# Both of those resources can be crucial for the performance of mobile applications.
+#
+
+import torch
+import torchvision
+
+supported_qengines = torch.backends.quantized.supported_engines
+print(supported_qengines)
+model = torchvision.models.quantization.__dict__['mobilenet_v2'](pretrained=True, quantize=True)
+torch.backends.quantized.engine='qnnpack'
+model.eval()
+script_model = torch.jit.script(model)
+x = torch.rand(1, 3 , 224, 224)
+y = script_model(x)
+torch.jit.save(script_model, 'mobilenetV2_quantized.pt') 
+
+
+######################################################################
+# Android API Recommendations
+# ---------------------------
+#
+# 4. Android. Reusing tensors for forward.
+#
+# Memory is a critical resource for android performance, especially on old devices.
+# Tensors can need a significant amount of memory.
+# For example, standard computer vision tensor contains 1*3*224*224 elements,
+# assuming that data type is float and will need 588Kb of memory.
+
+FloatBuffer buffer = Tensor.allocateFloatBuffer(1*3*224*224);
+Tensor tensor = Tensor.fromBlob(buffer, new long[]{1, 3, 224, 224});
+
+
+# Here we allocate native memory as ``java.nio.FloatBuffer`` and creating ``org.pytorch.Tensor`` which storage will be pointing to the memory of the allocated buffer.
+#
+# For most of the use cases, we do not do model forward only once, repeating it with some frequency or as fast as possible.
+#
+# If we are doing new memory allocation for every module forward - that will be suboptimal.
+# Instead of this, we can reuse the same memory that we allocated on the previous step, fill it with new data, and run module forward again on the same tensor object.
+#
+# You can check how it looks in code in [pytorch android application example](https://github.com/pytorch/android-demo-app/blob/master/PyTorchDemoApp/app/src/main/java/org/pytorch/demo/vision/ImageClassificationActivity.java#L174).
+
+protected AnalysisResult analyzeImage(ImageProxy image, int rotationDegrees) {
+  if (mModule == null) {
+    mModule = Module.load(moduleFileAbsoluteFilePath);
+    mInputTensorBuffer =
+    Tensor.allocateFloatBuffer(3 * 224 * 224);
+    mInputTensor = Tensor.fromBlob(mInputTensorBuffer, new long[]{1, 3, 224, 224});
+  }
+
+  TensorImageUtils.imageYUV420CenterCropToFloatBuffer(
+      image.getImage(), rotationDegrees,
+      224, 224,
+      TensorImageUtils.TORCHVISION_NORM_MEAN_RGB,
+      TensorImageUtils.TORCHVISION_NORM_STD_RGB,
+      mInputTensorBuffer, 0);
+
+  Tensor outputTensor = mModule.forward(IValue.from(mInputTensor)).toTensor();
+}
+
+
+# Member fields ``mModule``, ``mInputTensorBuffer`` and ``mInputTensor`` are initialized only once
+# and buffer is refilled using ``org.pytorch.torchvision.TensorImageUtils.imageYUV420CenterCropToFloatBuffer``.
+#
+#
+# Benchmarking
+# ------------
+#
+# The best way to benchmark (to check if optimizations helped your use case) - to measure your particular use case that you want to optimize, as performance behavior can vary in different environments.
+#
+# PyTorch distribution provides a way to benchmark naked binary that runs the model forward,
+# this approach can give more stable measurements rather than testing inside the application.
+# Android
+# -------
+# For this you first need to build benchmark binary:
+
+<from-your-root-pytorch-dir>
+rm -rf build_android
+BUILD_PYTORCH_MOBILE=1 ANDROID_ABI=arm64-v8a ./scripts/build_android.sh -DBUILD_BINARY=ON
+
+
+# You should have arm64 binary at: ``build_android/bin/speed_benchmark_torch``.
+# This binary takes --model=<path-to-model>, --input_dim="1,3,224,224 as dimension information for the input and --input_type="float" as the type of the input as arguments.
+#
+# Once you have your android device connected, push speedbenchark_torch binary and your model to the phone:
+
+adb push <speedbenchmark-torch> /data/local/tmp
+adb push <path-to-scripted-model> /data/local/tmp
+
+
+# Now we are ready to benchmark your model:
+#
+
+adb shell "/data/local/tmp/speed_benchmark_torch --model="/data/local/tmp/model.pt" --input_dims="1,3,224,224" --input_type="float"
+----- output -----
+Starting benchmark.
+Running warmup runs.
+Main runs.
+Main run finished. Microseconds per iter: 121318. Iters per second: 8.24281
+
+