torch.jit.optimize_for_inference

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

Выполняет набор оптимизационных проходов для оптимизации модели в целях вывода. Если модель еще не заморожена, optimize_for_inference вызовет torch.jit.freeze автоматически.

Помимо общих оптимизаций, которые должны ускорить вашу модель независимо от среды, подготовка к выводу также внедрит настройки, специфичные для построения, такие как наличие CUDNN или MKLDNN, и, возможно, в будущем произведет преобразования, ускоряющие работу на одной машине, но замедляющие ее на другой. Соответственно, сериализация не реализована после вызова optimize_for_inference и не гарантируется.

Это все еще прототип, и может потенциально замедлить вашу модель. Основные примеры использования, на которые до сих пор были нацелены, — это модели компьютерного зрения на процессоре и в меньшей степени на графическом процессоре.

Пример (оптимизация модуля с 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