torch.jit.optimize_for_inference

torch.jit.optimize_for_inference(mod, other_methods=None)[source][source]

对模型进行一系列优化操作，以优化推理目的。

如果模型尚未冻结，则 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

返回类型:

脚本模块