Функция torch.inference_mode в PyTorch

Pytorch torch 参考手册Справочное руководство Pytorch torch

torch.inference_modeявляется контекстным менеджером в PyTorch, используемым для режима инференса. Он болееtorch.no_gradстрогий: не только отключает вычисление градиентов, но и отключает все функции отслеживания механизма autograd.

Это, во время вывода модели, болееno_gradэффективно и может дополнительно снизить использование памяти и повысить скорость вывода.

Определение функции

torch.inference_mode(mode=True)

Parameters:

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

Возвращаемое значение:

  • Возвращает контекстный менеджер, в котором градиенты и autograd отключены.

Примеры использования

Пример 1: Базовое использование

Пример

import torch

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

# В контексте inference_mode
with torch.inference_mode():
    y = x * 2
    print("В inference_mode:", y.requires_grad)

# В контексте no_grad
with torch.no_grad():
    z = x * 2
    print("In no_grad:", z.requires_grad)

Вывод:

在 inference_mode 中: False
在 no_grad 中: False

Пример 2: Сравнение no_grad и inference_mode

Пример

import torch

# Создать тензор
x = torch.randn(100, 100)

# В inference_mode
with torch.inference_mode():
    # Выполнено много вычислений
    for _ in range(10):
        x = torch.mm(x, x)

    # Даже после завершения вычислений тензоры в контексте нельзя использовать для обратного распространения ошибки
    result = x.sum()

    # Проверить, можно ли его преобразовать в тензор, требующий градиентов
    print("В inference_mode:", result.is_leaf)

# Выполните те же вычисления в no_grad
x2 = torch.randn(100, 100)
with torch.no_grad():
    for _ in range(10):
        x2 = torch.mm(x2, x2)
    result2 = x2.sum()
    print("В no_grad:", result2.is_leaf)

Вывод:

在 inference_mode 中: False
在 no_grad 中: True

Пример 3: Инференс модели

Пример

import torch
import torch.nn as nn

# Определим простую модель
model = nn.Sequential(
    nn.Linear(10, 20),
    nn.ReLU(),
    nn.Linear(20, 5)
)

model.eval()

# Создайте входные данные
x = torch.randn(100, 10)

# Используйте inference_mode для инференса
with torch.inference_mode():
    output = model(x)
    print("Форма выхода:", output.shape)
    print("requires_grad выхода:", output.requires_grad)

# Можно также использовать декоратор
@torch.inference_mode()
def predict(x):
    return model(x)

result = predict(x)
print("Метод с декоратором - форма вывода:", result.shape)

Вывод:

输出形状: torch.Size([100, 5])
输出 requires_grad: False
装饰器方式 - 输出形状: torch.Size([100, 5])

Пример 4: Сравнение оптимизации памяти

Пример

import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Linear(1000, 1000),
    nn.ReLU(),
    nn.Linear(1000, 1000),
    nn.ReLU(),
    nn.Linear(1000, 10)
)

# Проверка использования памяти в различных режимах
x = torch.randn(50, 1000)

print("Без какого-либо контекстного менеджера:")
_ = model(x)

print("\nИспользование no_grad:")
with torch.no_grad():
    _ = model(x)

print("\nИспользование inference_mode:")
with torch.inference_mode():
    _ = model(x)

Использованиеinference_modeможет дополнительно оптимизировать память, поскольку полностью отключает движок autograd.

Связанные функции

  • torch.no_grad(): Отключает вычисление градиентов, но всё ещё сохраняет часть функциональности autograd.
  • torch.enable_grad(): Включает вычисление градиентов.
  • torch.is_inference_mode_enabled(): Проверяет, включён ли режим inference.

Примечания

  • inference_modeСравнениеno_gradстроже, отключает больше функциональности.
  • Вinference_modeТензоры, созданные в нём, помечаются как нелистовые узлы и не могут использоваться для обратного распространения.
  • Рекомендуется использовать во время вывода и оценки моделиinference_modeдля получения наилучшей производительности.
  • inference_modeНе может бытьno_gradвложенным.

Pytorch torch 参考手册Справочное руководство Pytorch torch

Другие расширения