torch.utils.mobile_optimizer

警告

PyTorch Mobile 已不再积极支持。请关注 ExecuTorch，它是 PyTorch 全新的设备端推理库。您还可以查阅关于 XNNPACK 和 Vulkan 委托的文档。

Torch Mobile 支持 torch.utils.mobile_optimizer.optimize_for_mobile 工具，用于在 eval 模式下对模块运行一系列优化。该方法接受以下参数：一个 torch.jit.ScriptModule 对象、一个优化阻止列表集合、一个要保留的方法列表以及一个后端。

对于 CPU 后端，默认情况下，如果优化阻止列表为 None 或为空，optimize_for_mobile 将运行以下优化：

• Conv2D + BatchNorm 融合 (阻止列表选项 mobile_optimizer.MobileOptimizerType.CONV_BN_FUSION)：此优化过程会将该模块及其所有子模块的 forward 方法中的 Conv2d-BatchNorm2d 折叠（融合）到 Conv2d 中。Conv2d 的权重和偏差会相应更新。

• 插入并折叠预打包算子 (阻止列表选项 mobile_optimizer.MobileOptimizerType.INSERT_FOLD_PREPACK_OPS)：此优化过程会重写图，将 2D 卷积和线性算子替换为其预打包的对应算子。预打包算子是有状态的，它们需要创建一些状态，例如权重预打包，并在算子执行期间使用此状态（即预打包权重）。XNNPACK 就是提供预打包算子的后端之一，其内核针对移动平台（如 ARM CPU）进行了优化。权重的预打包可以实现高效的内存访问，从而加快内核执行速度。目前，optimize_for_mobile 过程会将图重写，用以下两部分替换 Conv2D/Linear：1) 为 XNNPACK conv2d/linear 算子预打包权重的算子；2) 接受预打包权重和激活作为输入并生成输出激活的算子。由于步骤 1 只需要执行一次，我们对权重预打包进行折叠，使其仅在模型加载时执行一次。optimize_for_mobile 的此过程执行步骤 1 和 2，然后折叠（即删除）权重预打包算子。

• ReLU/Hardtanh 融合：XNNPACK 算子支持钳位融合。这意味着输出激活的钳位作为内核的一部分完成，包括 2D 卷积和线性算子内核。因此，钳位实际上是免费的。因此，任何可以表示为钳位算子的算子，例如 ReLU 或 hardtanh，都可以与 XNNPACK 中之前的 Conv2D 或 linear 算子融合。此过程通过查找前一个过程写入的跟随在 XNNPACK Conv2D/linear 算子后面的 ReLU/hardtanh 算子，并将其融合在一起，从而重写图。

• Dropout 删除 (阻止列表选项 mobile_optimizer.MobileOptimizerType.REMOVE_DROPOUT)：当训练标志为 false 时，此优化过程会从此模块中删除 dropout 和 dropout_ 节点。

• 卷积打包参数提升 (阻止列表选项 mobile_optimizer.MobileOptimizerType.HOIST_CONV_PACKED_PARAMS)：此优化过程会将卷积打包参数移动到根模块，以便可以删除卷积结构体。这可以在不影响数值的情况下减小模型大小。

• Add/ReLU 融合 (阻止列表选项 mobile_optimizer.MobileOptimizerType.FUSE_ADD_RELU)：此过程会查找跟随在 add 算子后面的 relu 算子实例，并将它们融合到一个 add_relu 中。

对于 Vulkan 后端，默认情况下，如果优化阻止列表为 None 或为空，optimize_for_mobile 将运行以下优化：

• 自动 GPU 传输 (阻止列表选项 mobile_optimizer.MobileOptimizerType.VULKAN_AUTOMATIC_GPU_TRANSFER)：此优化过程会重写图，以便输入和输出数据在 GPU 上的移动成为模型的一部分。

optimize_for_mobile 还会调用 freeze_module 过程，该过程只保留 forward 方法。如果您需要保留其他方法，请将其添加到保留方法列表并传递给该方法。

torch.utils.mobile_optimizer.optimize_for_mobile(script_module, optimization_blocklist=None, preserved_methods=None, backend='CPU')

优化 torch script 模块以用于移动部署。

参数

• script_module (ScriptModule) – 一个类型为 ScriptModule 的 torch script 模块实例。

• optimization_blocklist (Optional[set[torch._C._MobileOptimizerType]]) – 一个类型为 MobileOptimizerType 的集合。如果未传递该集合，优化方法将运行所有优化过程；否则，优化方法将运行不包含在 optimization_blocklist 中的优化过程。

• preserved_methods (Optional[list[~AnyStr]]) – 在调用 freeze_module 过程时需要保留的方法列表

• backend (str) – 用于运行结果模型的设备类型（“CPU”（默认）、“Vulkan”或“Metal”）。

返回

一个新的优化后的 torch script 模块

返回类型

RecursiveScriptModule