Функция inference_mode

Функция inference_mode в PyTorch создаёт контекстный менеджер, который включает режим вывода (inference mode). Этот режим является более строгой версией режима no_grad: он полностью отключает отслеживание градиентов, а также отключает проверки и механизмы, которые могут замедлять выполнение кода. Режим вывода предназначен исключительно для инференса моделей и значительно повышает производительность по сравнению с другими режимами.

Синтаксис

torch.inference_mode([mode])

Необязательный параметр mode имеет тип bool. Если передано True (значение по умолчанию), режим вывода включается; если False - отключается.

Пример

Давайте рассмотрим базовое использование режима вывода для инференса модели:

import torch

t = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)

with torch.inference_mode():
	res = t * 2
	print(res)
	print(res.requires_grad)

Результат выполнения кода:

tensor([2., 4., 6.])
False

Как видно из примера, созданный тензор не отслеживает градиенты, несмотря на то, что исходный тензор имел requires_grad=True.

Пример

Сравним производительность режимов inference_mode и no_grad при выполнении операций с тензорами:

import torch
import time

t = torch.randn(1000, 1000)

# With inference_mode
start = time.time()
with torch.inference_mode():
	for _ in range(100):
		res = t * 2
end = time.time()
print(f"inference_mode: {end - start:.4f}s")

# With no_grad
start = time.time()
with torch.no_grad():
	for _ in range(100):
		res = t * 2
end = time.time()
print(f"no_grad: {end - start:.4f}s")

Результат выполнения кода:

"inference_mode: 0.0023s"
"no_grad: 0.0041s"

Режим inference_mode работает быстрее благодаря отключению дополнительных проверок и оптимизациям, специфичным для инференса.

Пример

Покажем, что в режиме вывода невозможно выполнить операции, требующие вычисления градиентов, такие как backward:

import torch

t = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)

with torch.inference_mode():
	res = t * 2
	try:
		res.backward()
	except RuntimeError as e:
		print(e)

Результат выполнения кода:

"Tensors in inference mode cannot be used for autograd tracking"

Пример

Используем inference_mode для инференса нейронной сети:

import torch

torch.manual_seed(0)

# Создаём простую модель
model = torch.nn.Linear(10, 5)

# Входные данные
x = torch.randn(3, 10)

# Инференс в режиме вывода
with torch.inference_mode():
	output = model(x)
	print(output)
	print(output.requires_grad)

Результат выполнения кода:

tensor([
	[ 0.0415,  0.0596, -0.3817,  0.4716, -0.1115],
	[-0.2252,  0.1745, -0.0313, -0.1559, -0.0961],
	[-0.2257,  0.0136, -0.0734,  0.1760,  0.0562],
])
False

Модель работает в режиме инференса, не отслеживая градиенты, что экономит память и ускоряет вычисления.

Пример

Используем параметр mode для условного включения режима вывода:

import torch

t = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)

is_inference = True

with torch.inference_mode(mode=is_inference):
	res = t * 2
	print(res.requires_grad)

is_inference = False

with torch.inference_mode(mode=is_inference):
	res = t * 2
	print(res.requires_grad)

Результат выполнения кода:

False
True

Смотрите также

функцию no_grad,
которая отключает вычисление градиентов на время выполнения блока кода
функцию enable_grad,
которая включает вычисление градиентов для блока кода
функцию set_grad_enabled,
которая устанавливает режим вычисления градиентов глобально
функцию is_grad_enabled,
которая проверяет, включён ли режим вычисления градиентов в текущем контексте