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上的视觉模型，尽管在GPU上的应用较少。

示例（优化一个包含卷积和批量归一化的模块）：

```python
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)
# 如果使用MKLDNN构建，卷积将使用MKLDNN权重运行
assert "MKLDNN" in frozen_mod.graph
```

Return type: 脚本模块