Функция PyTorch torch.set_grad_enabled

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

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

Это очень полезно, когда нужно динамически управлять вычислением градиента в зависимости от условий.

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

torch.set_grad_enabled(mode)

Parameters:

mode(bool): Если оноTrue, включить вычисление градиента; если оноFalse, отключить вычисление градиента.

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

Возвращает контекстный менеджер, который можно использовать вwithинструкциях.

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

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

Пример

import torch

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

# Отключить вычисление градиента
torch.set_grad_enabled(False)
y1 = x * 2
print("После отключения градиентов:", y1.requires_grad)

# Включить вычисление градиента
torch.set_grad_enabled(True)
y2 = x * 2
print("После включения градиентов:", y2.requires_grad)

Вывод:

禁用梯度后: False
启用梯度后: True

Пример 2: Использование в качестве контекстного менеджера

Пример

import torch

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

# Использовать контекстный менеджер
with torch.set_grad_enabled(False):
y1 = x * 2
print("Внутри контекстного менеджера:", y1.requires_grad)

y2 = x * 2
print("Вне контекстного менеджера:", y2.requires_grad)

Вывод:

在上下文管理器内: False
在上下文管理器外: True

Пример 3: Динамическое управление обучением и инференсом

Пример

import torch
import torch.nn as nn

model = nn.Linear(10, 2)

def forward_pass(x, training=True):
"""Управление градиентами на основе параметра training"""
with torch.set_grad_enabled(training):
output = model(x)
print(f"training={training}, requires_grad={output.requires_grad}")
return output

# Режим обучения
x = torch.randn(5, 10)
forward_pass(x, training=True)

# Режим инференса
forward_pass(x, training=False)

Вывод:

training=True, requires_grad=True
training=False, requires_grad=False

Пример 4: Сохранение и восстановление состояния градиента

Пример

import torch

# Начальное состояние
print("Начальное состояние:", torch.is_grad_enabled())

# Создать контекстный менеджер, который возвращает в исходное состояние
old = torch.is_grad_enabled()

# Временно отключить градиенты
with torch.set_grad_enabled(False):
print("Внутри:", torch.is_grad_enabled())

# Автоматически восстановить (но здесь нам нужно восстановить вручную)
torch.set_grad_enabled(old)
print("После восстановления:", torch.is_grad_enabled())

Вывод:

初始状态: True
在内部: False
恢复后: True

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

torch.no_grad(): Контекстный менеджер, который отключает вычисление градиентов.
torch.enable_grad(): Контекстный менеджер, который включает вычисление градиентов.
torch.is_grad_enabled(): Проверяет, включено ли вычисление градиентов в данный момент.

Примечания

set_grad_enabledЕё можно вызывать напрямую как функцию для изменения глобального состояния или использовать в качестве контекстного менеджера.
При вызове как функции нужно вручную восстановить исходное состояние, иначе это повлияет на последующий код.
Рекомендуется использовать подход с контекстным менеджером, чтобы гарантировать правильное восстановление состояния.
Учитывайте влияние на глобальное состояние и используйте её с осторожностью в сложном коде.

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

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