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