• 文档 >
  • torch.utils.mobile_optimizer
快捷方式

torch.utils.mobile_optimizer

警告

PyTorch Mobile 现已停止积极维护。请查看 ExecuTorch,PyTorch 的全新设备端推理库。您还可以查看有关 XNNPACKVulkan 代理的文档。

Torch mobile 支持 torch.utils.mobile_optimizer.optimize_for_mobile 实用程序,用于对处于评估模式的模块运行一系列优化流程。此方法采用以下参数: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 流程重写图以使用 1) 为 XNNPACK conv2d/linear 操作预打包权重的操作和 2) 以预打包权重和激活作为输入并生成输出激活的操作来替换 Conv2D/Linear。由于 1 只需要执行一次,因此我们将权重预打包折叠起来,使其仅在模型加载时执行一次。此流程的 optimize_for_mobile 执行 1 和 2,然后折叠(即删除)权重预打包操作。

  • **ReLU/Hardtanh 融合**:XNNPACK 操作支持钳位融合。也就是说,输出激活的钳位作为内核的一部分完成,包括 2D 卷积和线性操作内核。因此,钳位实际上是免费的。因此,任何可以表示为钳位操作的操作,例如 ReLUhardtanh,都可以与 XNNPACK 中的先前 Conv2Dlinear 操作融合。此流程通过查找遵循 XNNPACK Conv2D/linear 操作(由上一流程编写)的 ReLU/hardtanh 操作并将其融合在一起重写图。

  • **删除 Dropout**(黑名单选项 mobile_optimizer.MobileOptimizerType.REMOVE_DROPOUT):此优化流程在训练为假时从此模块中删除 dropoutdropout_ 节点。

  • **卷积打包参数提升**(黑名单选项 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 和从 GPU 移动到 CPU 成为模型的一部分。

optimize_for_mobile 还将调用 freeze_module 流程,该流程仅保留 forward 方法。如果您有其他需要保留的方法,请将它们添加到保留方法列表中并传递给方法。

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

优化用于移动部署的 torch 脚本模块。

参数
  • script_module (ScriptModule) – 一个 torch 脚本模块的实例,类型为 ScriptModule。

  • optimization_blocklist (Optional[Set[_MobileOptimizerType]]) – 一个类型为 MobileOptimizerType 的集合。当未传递集合时,优化方法将运行所有优化器传递;否则,优化器方法将运行不在 optimization_blocklist 中的优化传递。

  • preserved_methods (Optional[List]) – 当调用 freeze_module 传递时需要保留的方法列表。

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

返回值

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

返回类型

RecursiveScriptModule

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取针对初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源