Python PyTorch optimize_for_inference用法及代碼示例

本文簡要介紹python語言中 torch.jit.optimize_for_inference 的用法。

用法:

torch.jit.optimize_for_inference(mod)

執行一組優化傳遞以優化模型以進行推理。如果模型尚未凍結，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=.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