Función PyTorch torch.set_grad_enabled

Pytorch torch 参考手册Manual de referencia de Pytorch torch

torch.set_grad_enabledEs una función en PyTorch que se utiliza para configurar globalmente si el cálculo de gradientes está habilitado. Puede habilitar o deshabilitar dinámicamente el cálculo de gradientes; a diferencia de un administrador de contexto, es una función que puede cambiar el estado global.

Esto es muy útil cuando se necesita controlar dinámicamente el cálculo de gradientes según condiciones.

Definición de la función

torch.set_grad_enabled(mode)

Parámetros:

  • mode(bool): Si esTrue, se habilita el cálculo de gradientes; si esFalse, se deshabilita el cálculo de gradientes.

Valor de retorno:

  • Devuelve un administrador de contexto que se puede utilizar parawithdeclaraciones.

Ejemplo de uso

Ejemplo 1: Uso básico

Ejemplo

import torch

# Crear un tensor que requiere gradiente
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)

# Deshabilitar el cálculo de gradientes
torch.set_grad_enabled(False)
y1 = x * 2
print("Después de deshabilitar el gradiente:", y1.requires_grad)

# Habilitar el cálculo de gradientes
torch.set_grad_enabled(True)
y2 = x * 2
print("Después de habilitar el gradiente:", y2.requires_grad)

El resultado de salida es:

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

Ejemplo 2: Uso como administrador de contexto

Ejemplo

import torch

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

# Usar el administrador de contexto
with torch.set_grad_enabled(False):
    y1 = x * 2
    print("Dentro del administrador de contexto:", y1.requires_grad)

y2 = x * 2
print("Fuera del administrador de contexto:", y2.requires_grad)

El resultado de salida es:

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

Ejemplo 3: Control dinámico de entrenamiento e inferencia

Ejemplo

import torch
import torch.nn as nn

model = nn.Linear(10, 2)

def forward_pass(x, training=True):
    """Controlar el gradiente según el parámetro training"""
    with torch.set_grad_enabled(training):
        output = model(x)
        print(f"training={training}, requires_grad={output.requires_grad}")
    return output

# Modo de entrenamiento
x = torch.randn(5, 10)
forward_pass(x, training=True)

# Modo de inferencia
forward_pass(x, training=False)

El resultado de salida es:

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

Ejemplo 4: Guardar y restaurar el estado del gradiente

Ejemplo

import torch

# Estado inicial
print("Estado inicial:", torch.is_grad_enabled())

# Crear un administrador de contexto que devuelva el estado original
old = torch.is_grad_enabled()

# Deshabilitar temporalmente el gradiente
with torch.set_grad_enabled(False):
    print("En el interior:", torch.is_grad_enabled())

# Restauración automática (pero aquí se necesita restaurar manualmente)
torch.set_grad_enabled(old)
print("Después de restaurar:", torch.is_grad_enabled())

El resultado de salida es:

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

Funciones relacionadas

  • torch.no_grad(): Administrador de contexto que deshabilita el cálculo de gradientes.
  • torch.enable_grad(): Administrador de contexto que habilita el cálculo de gradientes.
  • torch.is_grad_enabled(): Comprueba si el cálculo de gradientes está habilitado actualmente.

Notas

  • set_grad_enabledSe puede llamar directamente como función para cambiar el estado global, o usarse como administrador de contexto.
  • Cuando se llama como función, es necesario restaurar manualmente el estado original; de lo contrario, afectará al código posterior.
  • Se recomienda usar el administrador de contexto para asegurar que el estado se restaure correctamente.
  • Preste atención al impacto del estado global; úselo con cuidado en código complejo.

Pytorch torch 参考手册Manual de referencia de Pytorch torch

Otras extensiones