Fuse 模块配方¶

创建日期：2020年10月26日 | 最后更新：2021年7月26日 | 最后验证：2024年11月5日

本教程演示了如何将一系列PyTorch模块融合成一个单一模块，以及如何进行性能测试以比较融合模型与其非融合版本。

介绍¶

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

先决条件¶

PyTorch 1.6.0 或 1.7.0

步骤¶

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

1. 定义示例模型¶

使用在PyTorch Mobile Performance Recipes中定义的相同示例模型：

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

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

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
bn, relu

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

了解更多¶

请参阅此处以获取torch.quantization.fuse_modules的官方文档。