torch.jit.optimize_for_inference

torch.jit.optimize_for_inference(mod, other_methods=None)

執行一系列最佳化過程，以最佳化模型用於推理。

如果模型尚未凍結，optimize_for_inference 將自動呼叫 torch.jit.freeze。

除了通用的、無論在哪種環境下都能加速模型的最佳化之外，為推理做準備還將烘焙特定於構建的設定，例如 CUDNN 或 MKLDNN 的存在，並且將來可能會進行在一臺機器上加速但在另一臺機器上減速的轉換。因此，在呼叫 optimize_for_inference 後不實現序列化，也不保證序列化。

這仍處於原型階段，並有可能減慢您的模型。到目前為止，主要的用例是 CPU 上的視覺模型，以及在較小程度上是 GPU 上的視覺模型。

示例（最佳化具有 Conv->Batchnorm 的模組）

import torch

in_channels, out_channels = 3, 32
conv = torch.nn.Conv2d(
    in_channels, out_channels, kernel_size=3, stride=2, bias=True
)
bn = torch.nn.BatchNorm2d(out_channels, eps=0.001)
mod = torch.nn.Sequential(conv, bn)
frozen_mod = torch.jit.optimize_for_inference(torch.jit.script(mod.eval()))
assert "batch_norm" not in str(frozen_mod.graph)
# if built with MKLDNN, convolution will be run with MKLDNN weights
assert "MKLDNN" in frozen_mod.graph

返回型別

ScriptModule