融合模块配方¶

本配方演示了如何将 PyTorch 模块列表融合成单个模块，以及如何进行性能测试以比较融合模型与其未融合版本。

简介¶

在将量化应用于模型以减小其大小和内存占用量之前（有关量化的详细信息，请参阅量化配方），可以先将模型中的模块列表融合成单个模块。融合是可选的，但它可以节省内存访问，使模型运行速度更快，并提高其准确性。

先决条件¶

PyTorch 1.6.0 或 1.7.0

步骤¶

按照以下步骤融合示例模型，对其进行量化、脚本化、优化以用于移动设备、保存并使用 Android 基准测试工具进行测试。

1. 定义示例模型¶

使用 PyTorch 移动性能配方中定义的相同示例模型。

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

2. 生成带有和不带有 fuse_modules 的两个模型¶

在上面模型定义下方添加以下代码并运行脚本。

model = AnnotatedConvBnReLUModel()
print(model)

def prepare_save(model, fused):
    model.qconfig = torch.quantization.get_default_qconfig('qnnpack')
    torch.quantization.prepare(model, inplace=True)
    torch.quantization.convert(model, inplace=True)
    torchscript_model = torch.jit.script(model)
    torchscript_model_optimized = optimize_for_mobile(torchscript_model)
    torch.jit.save(torchscript_model_optimized, "model.pt" if not fused else "model_fused.pt")

prepare_save(model, False)

model = AnnotatedConvBnReLUModel()
model_fused = torch.quantization.fuse_modules(model, [['bn', 'relu']], inplace=False)
print(model_fused)

prepare_save(model_fused, True)

原始模型及其融合版本的图形将按如下所示打印。

AnnotatedConvBnReLUModel(
  (conv): Conv2d(3, 5, kernel_size=(3, 3), stride=(1, 1), bias=False)
  (bn): BatchNorm2d(5, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (relu): ReLU(inplace=True)
  (quant): QuantStub()
  (dequant): DeQuantStub()
)

AnnotatedConvBnReLUModel(
  (conv): Conv2d(3, 5, kernel_size=(3, 3), stride=(1, 1), bias=False)
  (bn): BNReLU2d(
    (0): BatchNorm2d(5, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (1): ReLU(inplace=True)
  )
  (relu): Identity()
  (quant): QuantStub()
  (dequant): DeQuantStub()
)

在第二个融合模型输出中，列表中的第一个项目 bn 被融合模块替换，其余模块（本例中为 relu）被标识替换。此外，还生成了模型 model.pt 和 model_fused.pt 的未融合版本和融合版本。

3. 构建 Android 基准测试工具¶

获取 PyTorch 源码并按照以下步骤构建 Android 基准测试工具

git clone --recursive https://github.com/pytorch/pytorch
cd pytorch
git submodule update --init --recursive
BUILD_PYTORCH_MOBILE=1 ANDROID_ABI=arm64-v8a ./scripts/build_android.sh -DBUILD_BINARY=ON

这将在 build_android/bin 文件夹中生成 Android 基准测试二进制文件 speed_benchmark_torch。

4. 测试比较融合和非融合模型¶

连接您的 Android 设备，然后复制 speed_benchmark_torch 和模型文件，并在其上运行基准测试工具

adb push build_android/bin/speed_benchmark_torch /data/local/tmp
adb push model.pt /data/local/tmp
adb push model_fused.pt /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"
adb shell "/data/local/tmp/speed_benchmark_torch --model=/data/local/tmp/model_fused.pt" --input_dims="1,3,224,224" --input_type="float"

最后两个命令的结果应如下所示

Main run finished. Microseconds per iter: 6189.07. Iters per second: 161.575

以及

Main run finished. Microseconds per iter: 6216.65. Iters per second: 160.858

对于此示例模型，融合模型和非融合模型之间没有太大性能差异。但是，可以使用类似的步骤来融合和准备真实的深度模型，并测试以查看性能改进。请记住，目前 torch.quantization.fuse_modules 仅融合以下模块序列

卷积层，批归一化层
卷积层，批归一化层，ReLU
卷积层，ReLU
线性层，ReLU
批归一化层，ReLU

如果向 fuse_modules 调用提供了任何其他序列列表，则将被简单忽略。

了解更多¶

有关 torch.quantization.fuse_modules 的官方文档，请参阅此处。