Pytorch 移动端性能优化技巧¶

警告

PyTorch Mobile 现已不再积极维护。请查看 ExecuTorch，PyTorch 的全新设备端推理库。您还可以详细了解量化、硬件加速（使用 hw 进行运算融合）和基准测试，这些内容在 ExecuTorch 的文档页面中有所介绍。

简介¶

性能（即延迟）对于大多数（如果不是全部）移动设备上 ML 模型推理的应用和用例至关重要。

目前，PyTorch 在 CPU 后端执行模型，其他硬件后端（如 GPU、DSP 和 NPU）尚不可用。

在本技巧中，您将学习

如何优化您的模型以帮助减少移动设备上的执行时间（提高性能，降低延迟）。
如何进行基准测试（以检查优化是否对您的用例有所帮助）。

模型准备¶

我们将从准备优化您的模型开始，以帮助减少移动设备上的执行时间（提高性能，降低延迟）。

设置¶

首先，我们需要使用 conda 或 pip 安装 PyTorch，版本至少为 1.5.0。

conda install pytorch torchvision -c pytorch

或者

pip install torch torchvision

编写你的模型

import torch
from torch.utils.mobile_optimizer import optimize_for_mobile

class AnnotatedConvBnReLUModel(torch.nn.Module):
    def __init__(self):
        super(AnnotatedConvBnReLUModel, self).__init__()
        self.conv = torch.nn.Conv2d(3, 5, 3, bias=False).to(dtype=torch.float)
        self.bn = torch.nn.BatchNorm2d(5).to(dtype=torch.float)
        self.relu = torch.nn.ReLU(inplace=True)
        self.quant = torch.quantization.QuantStub()
        self.dequant = torch.quantization.DeQuantStub()

    def forward(self, x):
        x = x.contiguous(memory_format=torch.channels_last)
        x = self.quant(x)
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        x = self.dequant(x)
        return x

model = AnnotatedConvBnReLUModel()

torch.quantization.QuantStub 和 torch.quantization.DeQuantStub() 是无操作存根，将在量化步骤中使用。

1. 使用 `torch.quantization.fuse_modules` 融合算子¶

不要误以为 `fuse_modules` 位于量化包中。它适用于所有 torch.nn.Module。

torch.quantization.fuse_modules 将一系列模块融合成单个模块。它仅融合以下模块序列：

卷积，批标准化
卷积，批标准化，ReLU
卷积，ReLU
线性，ReLU

此脚本将融合之前声明的模型中的卷积、批标准化和 ReLU。

torch.quantization.fuse_modules(model, [['conv', 'bn', 'relu']], inplace=True)

2. 量化您的模型¶

您可以在此教程中找到更多关于 PyTorch 量化的信息。

模型的量化不仅将计算转移到 int8，还减小了模型在磁盘上的大小。这种大小的减小有助于减少模型首次加载期间的磁盘读取操作，并减少 RAM 的使用量。这两种资源对于移动应用程序的性能都至关重要。此代码执行量化，使用存根作为模型校准函数，您可以在此处找到更多相关信息。

model.qconfig = torch.quantization.get_default_qconfig('qnnpack')
torch.quantization.prepare(model, inplace=True)
# Calibrate your model
def calibrate(model, calibration_data):
    # Your calibration code here
    return
calibrate(model, [])
torch.quantization.convert(model, inplace=True)

3. 使用 `torch.utils.mobile_optimizer`¶

Torch mobile_optimizer 包对脚本化的模型进行了一些优化，这将有助于卷积和线性操作。它以优化的格式预打包模型权重，并在上面融合操作和 ReLU（如果它是下一个操作）。

首先，我们对上一步得到的模型进行脚本化。

torchscript_model = torch.jit.script(model)

接下来，我们调用 optimize_for_mobile 并将模型保存到磁盘。

torchscript_model_optimized = optimize_for_mobile(torchscript_model)
torch.jit.save(torchscript_model_optimized, "model.pt")

4. 优先使用 Channels Last 张量内存格式¶

Channels Last (NHWC) 内存格式在 PyTorch 1.4.0 中引入。它仅支持四维张量。这种内存格式为大多数运算符（尤其是卷积）提供了更好的内存局部性。我们的测量结果表明，与默认的 Channels First (NCHW) 格式相比，MobileNetV2 模型的速度提高了 3 倍。

在编写此配方时，PyTorch Android Java API 不支持使用 Channels Last 内存格式的输入。但它可以在 TorchScript 模型级别使用，方法是为模型输入添加转换。

def forward(self, x):
    x = x.contiguous(memory_format=torch.channels_last)
    ...

如果您的输入已使用 Channels Last 内存格式，则此转换不会产生任何成本。之后，所有运算符都将保留 ChannelsLast 内存格式。

5. Android - 重用张量进行前向传播¶

此配方的一部分仅适用于 Android。

内存是 Android 性能的关键资源，尤其是在旧设备上。张量可能需要大量的内存。例如，标准计算机视觉张量包含 1*3*224*224 个元素，假设数据类型为浮点型，将需要 588Kb 的内存。

FloatBuffer buffer = Tensor.allocateFloatBuffer(1*3*224*224);
Tensor tensor = Tensor.fromBlob(buffer, new long[]{1, 3, 224, 224});

在这里，我们将本机内存分配为 java.nio.FloatBuffer 并创建 org.pytorch.Tensor，其存储将指向分配的缓冲区的内存。

在大多数情况下，我们不会只进行一次模型前向传播，而是会以一定的频率或尽可能快地重复进行。

如果我们对每个模块前向传播都进行新的内存分配，那将不是最佳方案。相反，我们可以重用在上一步分配的相同内存，用新数据填充它，然后在相同的张量对象上再次运行模块前向传播。

您可以在 pytorch android 应用程序示例中查看代码示例。

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

成员字段 mModule、mInputTensorBuffer 和 mInputTensor 只初始化一次，缓冲区使用 org.pytorch.torchvision.TensorImageUtils.imageYUV420CenterCropToFloatBuffer 重新填充。

6. 加载时间优化¶

自 Pytorch 1.13 起可用

PyTorch Mobile 还支持基于 FlatBuffer 的文件格式，该格式加载速度更快。基于 flatbuffer 和 pickle 的模型文件都可以使用相同的 _load_for_lite_interpreter（Python）或 ``_load_for_mobile``（C++）API 加载。

要使用 FlatBuffer 格式，而不是使用 model._save_for_lite_interpreter('path/to/file.ptl') 创建模型文件，您可以运行以下命令：

可以使用以下方式保存：

model._save_for_lite_interpreter('path/to/file.ptl', _use_flatbuffer=True)

额外的参数 _use_flatbuffer 会创建 FlatBuffer 文件而不是 zip 文件。创建的文件加载速度会更快。

例如，使用 ResNet-50 并运行以下脚本：

import torch
from torch.jit import mobile
import time
model = torch.hub.load('pytorch/vision:v0.10.0', 'deeplabv3_resnet50', pretrained=True)
model.eval()
jit_model = torch.jit.script(model)

jit_model._save_for_lite_interpreter('/tmp/jit_model.ptl')
jit_model._save_for_lite_interpreter('/tmp/jit_model.ff', _use_flatbuffer=True)

import timeit
print('Load ptl file:')
print(timeit.timeit('from torch.jit import mobile; mobile._load_for_lite_interpreter("/tmp/jit_model.ptl")',
                       number=20))
print('Load flatbuffer file:')
print(timeit.timeit('from torch.jit import mobile; mobile._load_for_lite_interpreter("/tmp/jit_model.ff")',
                       number=20))

您将得到以下结果：

Load ptl file:
0.5387594579999999
Load flatbuffer file:
0.038842832999999466

虽然实际移动设备上的加速效果会更小，但您仍然可以预期加载时间减少 3 倍到 6 倍。

### 避免使用基于 FlatBuffer 的移动模型的理由

但是，FlatBuffer 格式也有一些限制，您可能需要考虑：

它仅在 PyTorch 1.13 或更高版本中可用。因此，使用早期 PyTorch 版本编译的客户端设备可能无法加载它。
Flatbuffer 库对文件大小施加了 4GB 的限制。因此，它不适用于大型模型。

基准测试¶

基准测试（检查优化是否对您的用例有所帮助）的最佳方法是衡量您想要优化的特定用例，因为性能行为在不同的环境中可能会有所不同。

PyTorch 发行版提供了一种方法来基准测试运行模型前向传播的裸二进制文件，这种方法可以提供比在应用程序内部测试更稳定的测量结果。

Android - 基准测试设置¶

此配方的一部分仅适用于 Android。

为此，您首先需要构建基准测试二进制文件：

<from-your-root-pytorch-dir>
rm -rf build_android
BUILD_PYTORCH_MOBILE=1 ANDROID_ABI=arm64-v8a ./scripts/build_android.sh -DBUILD_BINARY=ON

您应该在 build_android/bin/speed_benchmark_torch 中拥有 arm64 二进制文件。此二进制文件将 --model=<path-to-model>、--input_dim="1,3,224,224" 作为输入的维度信息以及 --input_type="float" 作为输入的类型作为参数。

将您的 Android 设备连接后，将 `speedbenchark_torch` 二进制文件和您的模型推送到手机上。

adb push <speedbenchmark-torch> /data/local/tmp
adb push <path-to-scripted-model> /data/local/tmp

现在，我们已准备好对您的模型进行基准测试。

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

iOS - 基准测试设置¶

对于 iOS，我们将使用我们的 TestApp 作为基准测试工具。

首先，让我们将 optimize_for_mobile 方法应用于位于 TestApp/benchmark/trace_model.py 的 Python 脚本。只需修改如下代码即可。

import torch
import torchvision
from torch.utils.mobile_optimizer import optimize_for_mobile

model = torchvision.models.mobilenet_v2(pretrained=True)
model.eval()
example = torch.rand(1, 3, 224, 224)
traced_script_module = torch.jit.trace(model, example)
torchscript_model_optimized = optimize_for_mobile(traced_script_module)
torch.jit.save(torchscript_model_optimized, "model.pt")

现在，让我们运行 python trace_model.py。如果一切正常，我们应该能够在基准测试目录中生成优化的模型。

接下来，我们将从源代码构建 PyTorch 库。

BUILD_PYTORCH_MOBILE=1 IOS_ARCH=arm64 ./scripts/build_ios.sh

现在我们有了优化的模型和 PyTorch，是时候生成我们的 XCode 项目并进行基准测试了。为此，我们将使用一个 Ruby 脚本 - `setup.rb`，它负责设置 XCode 项目的繁重工作。

ruby setup.rb

现在打开 `TestApp.xcodeproj` 并插入您的 iPhone，您就可以开始了。以下是 iPhoneX 的示例结果：

TestApp[2121:722447] Main runs
TestApp[2121:722447] Main run finished. Milliseconds per iter: 28.767
TestApp[2121:722447] Iters per second: : 34.762
TestApp[2121:722447] Done.