目录

torch.utils.mobile_optimizer

警告

PyTorch Mobile 不再获得积极支持。请查看 ExecuTorch,这是 PyTorch 全新的设备端推理库。您还可以查阅有关 XNNPACKVulkan 代理的文档。

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 + 批归一化融合(阻止列表选项 mobile_optimizer.MobileOptimizerType.CONV_BN_FUSION):此优化过程将 Conv2d-BatchNorm2d 折叠到 Conv2d 中的 forward 方法中,并应用于该模块及其所有子模块。相应地更新 Conv2d 的权重和偏置。

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

  • ReLU/Hardtanh 融合: XNNPACK 操作支持融合裁剪。也就是说,输出激活的裁剪是在内核中完成的,包括二维卷积和线性操作内核。因此,裁剪实际上是免费的。因此,任何可以表示为裁剪操作的操作,例如 ReLUhardtanh,都可以与之前的 XNNPACK 中的 Conv2Dlinear 操作融合。此过程通过查找跟随 XNNPACK Conv2D/linear 操作(由上一个过程编写)的 ReLU/hardtanh 操作,并将它们融合在一起,从而重写图。

  • 丢弃移除 (阻止列表选项 mobile_optimizer.MobileOptimizerType.REMOVE_DROPOUT): 此优化过程在训练为 false 时从该模块中移除 dropoutdropout_ 节点。

  • 卷积打包参数提升(黑名单选项 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')[source]

为移动设备部署优化一个 torch 脚本模块。

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

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

  • 保留的方法 (可选[列表]) – 当调用freeze_module传递时需要保留的方法列表

  • 后端 (str) – 用于运行结果模型的设备类型(‘CPU’(默认), ‘Vulkan’ 或 ‘Metal’)。

Returns

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

Return type

RecursiveScriptModule

文档

访问 PyTorch 的全面开发人员文档

查看文档

教程

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

查看教程

资源

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

查看资源