torch.utils.mobile_optimizer¶
警告
此 API 处于测试阶段,可能会在不久的将来发生变化。
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
折叠为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 操作支持对钳位的融合。这意味着输出激活的钳位是在内核中完成的,包括二维卷积和线性操作内核。因此钳位实际上是免费的。因此,任何可以表示为钳位操作的操作,例如
ReLU
或hardtanh
,可以与 XNNPACK 中的前一个Conv2D
或linear
操作融合。此传递通过查找紧随 XNNPACKConv2D/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')[source]¶
针对移动部署优化 torch 脚本模块。
- 参数
script_module (ScriptModule) – torch 脚本模块实例,类型为 ScriptModule。
optimization_blocklist (可选[集合[_MobileOptimizerType]]) – 一个包含 MobileOptimizerType 类型元素的集合。当没有传递集合时,优化方法将运行所有优化器传递;否则,优化方法将运行不在 optimization_blocklist 中的优化器传递。
preserved_methods (可选[列表]) – 当调用 freeze_module 传递时需要保留的方法列表。
backend (字符串) – 用于运行结果模型的设备类型(‘CPU’(默认)、‘Vulkan’ 或 ‘Metal’)。
- 返回值
一个新的优化后的 torch 脚本模块
- 返回类型
RecursiveScriptModule