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=.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
- 返回类型
