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)：此最佳化過程會將該模組及其所有子模組的 forward 方法中的 Conv2d-BatchNorm2d 摺疊（融合）到 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 運算元支援鉗位融合。這意味著輸出啟用的鉗位作為核心的一部分完成，包括 2D 卷積和線性運算元核心。因此，鉗位實際上是免費的。因此，任何可以表示為鉗位運算元的運算元，例如 ReLU 或 hardtanh，都可以與 XNNPACK 中之前的 Conv2D 或 linear 運算元融合。此過程透過查詢前一個過程寫入的跟隨在 XNNPACK Conv2D/linear 運算元後面的 ReLU/hardtanh 運算元，並將其融合在一起，從而重寫圖。

• Dropout 刪除 (阻止列表選項 mobile_optimizer.MobileOptimizerType.REMOVE_DROPOUT)：當訓練標誌為 false 時，此最佳化過程會從此模組中刪除 dropout 和 dropout_ 節點。

• 卷積打包引數提升 (阻止列表選項 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[torch._C._MobileOptimizerType]]) – 一個型別為 MobileOptimizerType 的集合。如果未傳遞該集合，最佳化方法將執行所有最佳化過程；否則，最佳化方法將執行不包含在 optimization_blocklist 中的最佳化過程。

• preserved_methods (Optional[list[~AnyStr]]) – 在呼叫 freeze_module 過程時需要保留的方法列表

• backend (str) – 用於執行結果模型的裝置型別（“CPU”（預設）、“Vulkan”或“Metal”）。

返回

一個新的最佳化後的 torch script 模組

返回型別

RecursiveScriptModule