Функция PyTorch torch.no_grad

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

torch.no_grad— это контекстный менеджер в PyTorch, используемый для отключения вычисления градиентов. Вno_gradТензоры, созданные внутри блока, не вычисляют градиенты, что может значительно уменьшить потребление памяти и повысить скорость инференса.

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

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

torch.no_grad()

Parameters:

  • Нет параметров. Это контекстный менеджер.

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

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

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

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

Пример

import torch

# Создайте тензор, требующий вычисления градиентов
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)

# Выполните вычисления внутри контекста no_grad
with torch.no_grad():
    y = x * 2
    print("Внутри no_grad:", y.requires_grad)

# Вне no_grad
z = x * 2
print("Вне no_grad:", z.requires_grad)

Вывод:

在 no_grad 中: False
在 no_grad 外: True

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

Пример

import torch
import torch.nn as nn

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

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

# Используйте no_grad во время инференса
with torch.no_grad():
    output = model(x)
    print("Вывод:", output)
    # Вычисление градиента не требуется

# Или используйте декоратор
# @torch.no_grad()
# def predict(x):
# return model(x)

Вывод:

输出: tensor([[0.0920, 0.3557]])

Пример 3: Режим оценки

Пример

import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Linear(10, 5),
    nn.ReLU(),
    nn.Linear(5, 2)
)

# Переключиться в режим оценки
model.eval()

# Подготовка тестовых данных
test_input = torch.randn(5, 10)

# Отключить градиенты во время оценки
with torch.no_grad():
    predictions = model(test_input)
    print("Размер результата предсказания:", predictions.shape)

Вывод:

预测结果形状: torch.Size([5, 2])

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

Пример

import torch
import torch.nn as nn

model = nn.Linear(1000, 1000)

# Создать большой объем входных данных
inputs = [torch.randn(100, 1000) for _ in range(100)]

# Без no_grad (история градиентов будет записана)
print("Без no_grad:")
for inp in inputs[:5]:
    _ = model(inp)

# Используйте no_grad (не записывает историю градиентов)
print("С no_grad:")
with torch.no_grad():
    for inp in inputs[:5]:
        _ = model(inp)

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

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

  • torch.enable_grad(): Включает вычисление градиентов (в отличие отno_gradпротивоположного).
  • torch.set_grad_enabled(grad): Включает или отключает вычисление градиентов в зависимости от параметров.
  • torch.inference_mode(): Более строгий режим вывода, который отключает как вычисление градиентов, так и autograd.

Примечания

  • Во время вывода и оценки модели всегда используйтеtorch.no_grad()чтобы сэкономить память и повысить скорость.
  • Еслиno_gradпеременные, созданные внутри блока, должны использоваться вне блока, их нужно вручную скопировать наружу.
  • иmodel.eval()Используйте в сочетании для достижения наилучших результатов.

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

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