torch.jit.optimize_for_inference¶
- torch.jit.optimize_for_inference(mod, other_methods=None)[原始碼][原始碼]¶
執行一組最佳化流程，以最佳化模型以進行推論。
如果模型尚未凍結，optimize_for_inference 將自動調用 torch.jit.freeze。
除了通用最佳化（無論環境如何，都應該加速您的模型）之外，prepare for inference 還將包含建構特定的設定，例如 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=.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
- 傳回類型
