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）：此优化过程将 Conv2d-BatchNorm2d 折叠到此模块及其所有子模块的 forward 方法中的 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 操作支持 clamping 融合。也就是说，输出激活的 clamping 作为内核的一部分完成，包括 2D 卷积和线性操作内核。因此，clamping 实际上是免费的。因此，任何可以表示为 clamping 操作的操作，例如 ReLU 或 hardtanh，都可以与之前的 XNNPACK Conv2D 或 linear 操作融合。此过程通过查找跟在先前过程编写的 XNNPACK Conv2D/linear 操作之后的 ReLU/hardtanh 操作来重写图，并将它们融合在一起。

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

• Conv 打包参数提升（阻止列表选项 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[_MobileOptimizerType]]) – MobileOptimizerType 类型的集合。当未传递集合时，优化方法将运行所有优化过程；否则，优化方法将运行未包含在 optimization_blocklist 中的优化过程。

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

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

返回

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

返回类型

RecursiveScriptModule