Función PyTorch torch.enable_grad

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

torch.enable_grades un administrador de contexto en PyTorch para habilitar el cálculo de gradientes. Estorch.no_gradlo contrario, se utiliza para habilitar explícitamente el cálculo de gradientes en bloques de código que lo requieran.

Esto es muy útil cuando se necesita entrenar el modelo temporalmente en código de inferencia, o cuando se necesita habilitar el cálculo de gradientes en alguna área local.

Definición de la función

torch.enable_grad()

Parámetros:

  • Sin parámetros. Es un administrador de contexto.

Valor de retorno:

  • Devuelve un administrador de contexto en el que se habilita el cálculo de gradientes.

Ejemplos de uso

Ejemplo 1: Uso básico

Ejemplo

import torch

# Por defecto, en el contexto no_grad
with torch.no_grad():
    x = torch.tensor([1.0, 2.0, 3.0])
    print(Dentro de no_grad:, x.requires_grad)

    # Usar enable_grad para habilitar gradientes temporalmente
    with torch.enable_grad():
        y = x * 2
        print("En enable_grad:", y.requires_grad)

    # Después de salir, restaurar el estado no_grad
    z = x * 2
    print("Después de salir:", z.requires_grad)

El resultado es:

在 no_grad 中: False
在 enable_grad 中: True
退出后: False

Ejemplo 2: Uso mixto de entrenamiento e inferencia

Ejemplo

import torch
import torch.nn as nn

model = nn.Linear(10, 2)

# Modo de inferencia
with torch.no_grad():
    x = torch.randn(5, 10)
    output1 = model(x)
    print(Salida de inferencia:, output1.shape)

    # Si es necesario entrenar temporalmente algunos parámetros durante la inferencia
    with torch.enable_grad():
        # Crear un tensor que requiere gradiente para el cálculo
        temp_weight = torch.randn(10, 10, requires_grad=True)
        temp_output = torch.mm(x, temp_weight)
        print("Activar temporalmente el gradiente:", temp_output.requires_grad)

El resultado es:

推理输出: torch.Size([5, 2])
临时启用梯度: True

Ejemplo 3: Uso en un decorador

Ejemplo

import torch

@torch.enable_grad()
def train_step(x, y):
    """Simular pasos de entrenamiento"""
    # Esta función habilitará internamente el cálculo de gradientes
    loss = (x - y).sum()
    return loss

# Llamar en el contexto no_grad
with torch.no_grad():
    x = torch.tensor([1.0, 2.0, 3.0])
    y = torch.tensor([0.0, 0.0, 0.0])
    # El decorador asegura que el cálculo de gradiente esté habilitado
    loss = train_step(x, y)
    print("Loss:", loss)
    print(Requiere gradiente:, loss.requires_grad)

El resultado es:

Loss: tensor(6.)
需要梯度: True

Funciones relacionadas

  • torch.no_grad(): Deshabilita el cálculo de gradientes.
  • torch.set_grad_enabled(grad): Habilita o deshabilita el cálculo de gradientes según los parámetros.
  • torch.is_grad_enabled(): Comprueba si el cálculo de gradientes está habilitado actualmente.

Notas

  • enable_gradSe utiliza principalmente enno_gradcontextos para habilitar temporalmente el cálculo de gradientes.
  • Si se usa en un entorno donde los gradientes están habilitados globalmenteenable_gradno tendrá ningún efecto.
  • Se recomienda usarenable_gradel decorador para asegurar que los gradientes estén siempre habilitados dentro de la función.

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

Otras extensiones