快捷方式

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

ScriptModule

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取针对初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源