torch.utils.mobile_optimizer¶

警告

此API目前处于测试阶段，可能会在不久的将来发生变化。

Torch mobile 支持 torch.utils.mobile_optimizer.optimize_for_mobile 工具，用于在评估模式下对模块运行一系列优化传递。该方法接受以下参数：一个 torch.jit.ScriptModule 对象、一个黑名单优化集、一个保留方法列表和一个后端。

For CPU Backend, by default, if optimization blocklist is None or empty, optimize_for_mobile will run the following optimizations:

Conv2D + BatchNorm 融合（屏蔽选项 mobile_optimizer.MobileOptimizerType.CONV_BN_FUSION）：此优化过程将 Conv2d-BatchNorm2d 折叠为 Conv2d 在该模块及其所有子模块的 forward 方法中。Conv2d 的权重和偏置相应地更新。
插入和折叠预打包操作（阻止列表选项 mobile_optimizer.MobileOptimizerType.INSERT_FOLD_PREPACK_OPS）：此优化过程重写图形以将2D卷积和线性操作替换为其预打包的对应操作。预打包操作是有状态的操作，因为它们需要创建一些状态，例如权重预打包，并在操作执行期间使用此状态，即预打包的权重。XNNPACK是提供预打包操作的后端之一，其内核针对移动平台（如ARM CPU）进行了优化。权重的预打包可以实现高效的内存访问，从而加快内核执行速度。目前，optimize_for_mobile 过程重写图形以将 Conv2D/Linear 替换为1）预打包XNNPACK conv2d/线性操作权重的操作和2）接受预打包权重和激活作为输入并生成输出激活的操作。由于1只需要执行一次，我们将权重预打包折叠，使其仅在模型加载时执行一次。此 optimize_for_mobile 过程执行1和2，然后折叠，即删除权重预打包操作。
ReLU/Hardtanh 融合: XNNPACK 操作支持夹紧的融合。即输出激活的夹紧作为内核的一部分完成，包括对于 2D 卷积和线性操作内核。因此，夹紧实际上是免费的。因此，任何可以表示为夹紧操作的操作，例如 ReLU 或 hardtanh，都可以与之前的 Conv2D 或 linear 操作在 XNNPACK 中融合。此过程通过查找由前一过程编写的跟随 XNNPACK Conv2D/linear 操作的 ReLU/hardtanh 操作，并将它们融合在一起来重写图。
Dropout 移除（黑名单选项 mobile_optimizer.MobileOptimizerType.REMOVE_DROPOUT）：此优化过程在训练为假时从此模块中移除 dropout 和 dropout_ 节点。
卷积打包参数提升（阻止选项 mobile_optimizer.MobileOptimizerType.HOIST_CONV_PACKED_PARAMS）：此优化过程将卷积打包参数移动到根模块，以便可以删除卷积结构。这减少了模型大小，而不会影响数值。
添加/ReLU融合（排除选项 mobile_optimizer.MobileOptimizerType.FUSE_ADD_RELU）：此过程查找跟随在 add 操作之后的 relu 操作，并将它们融合成一个单独的 add_relu。

for Vulkan Backend, by default, if optimization blocklist is None or empty, optimize_for_mobile will run the following optimization:

自动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 脚本模块。

Parameters

script_module (ScriptModule) – 一个类型为ScriptModule的torch脚本模块实例。
optimization_blocklist (可选[集合[_MobileOptimizerType]]) – 一个包含MobileOptimizerType类型的集合。当未传递集合时，优化方法将运行所有优化过程；否则，优化方法将运行未包含在optimization_blocklist中的优化过程。
preserved_methods（可选[列表]）– 当调用freeze_module传递时需要保留的方法列表
backend (str) – 用于运行结果模型的设备类型（‘CPU’（默认），‘Vulkan’ 或 ‘Metal’）。

Returns

一个新的优化过的 torch 脚本模块

Return type

递归脚本模块