Función torch.no_grad de PyTorch

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

torch.no_grades un administrador de contexto en PyTorch utilizado para deshabilitar el cálculo de gradientes. En elno_gradbloque, los tensores creados no calcularán gradientes, lo que puede reducir significativamente el consumo de memoria y mejorar la velocidad de inferencia.

Esto es esencial en las etapas de inferencia y evaluación del modelo, ya que puede mejorar considerablemente el rendimiento y ahorrar memoria.

Definición de la función

torch.no_grad()

Parámetros:

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

Valor de retorno:

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

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)

# Calcular en el contexto de no_grad
with torch.no_grad():
    y = x * 2
    print("Dentro de no_grad:", y.requires_grad)

# Fuera de no_grad
z = x * 2
print("Fuera de no_grad:", z.requires_grad)

La salida es:

在 no_grad 中: False
在 no_grad 外: True

Ejemplo 2: Inferencia del modelo

Ejemplo

import torch
import torch.nn as nn

# Definir un modelo simple
model = nn.Linear(10, 2)

# Crear los datos de entrada
x = torch.randn(1, 10)

# Usar no_grad durante la inferencia
with torch.no_grad():
    output = model(x)
    print("Salida:", output)
    # No es necesario calcular gradientes

# O usar un decorador
# @torch.no_grad()
# def predict(x):
#     return model(x)

La salida es:

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

Ejemplo 3: Modo de evaluación

Ejemplo

import torch
import torch.nn as nn

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

# Cambiar al modo de evaluación
model.eval()

# Preparar los datos de prueba
test_input = torch.randn(5, 10)

# Deshabilitar gradientes durante la evaluación
with torch.no_grad():
    predictions = model(test_input)
    print("Forma de los resultados de predicción:", predictions.shape)

La salida es:

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

Ejemplo 4: Comparación del uso de memoria

Ejemplo

import torch
import torch.nn as nn

model = nn.Linear(1000, 1000)

# Crear una gran cantidad de entradas
inputs = [torch.randn(100, 1000) for _ in range(100)]

# Sin usar no_grad (se registrará el historial de gradientes)
print("Sin usar no_grad:")
for inp in inputs[:5]:
    _ = model(inp)

# Usar no_grad (no se registra el historial de gradientes)
print("Usando no_grad:")
with torch.no_grad():
    for inp in inputs[:5]:
        _ = model(inp)

Usarno_gradpuede reducir significativamente el consumo de memoria, ya que no es necesario almacenar información de gradientes para las variables intermedias.

Funciones relacionadas

  • torch.enable_grad(): Habilita el cálculo de gradientes (en oposición ano_gradlo contrario).
  • torch.set_grad_enabled(grad): Habilita o deshabilita el cálculo de gradientes según el parámetro.
  • torch.inference_mode(): Un modo de inferencia más estricto, que además deshabilita el cálculo de gradientes y autograd.

Notas importantes

  • Al realizar la inferencia y evaluación del modelo, utiliza siempretorch.no_grad()para ahorrar memoria y mejorar la velocidad.
  • Si las variables creadas en elno_gradbloque necesitan usarse fuera del bloque, deben copiarse manualmente.
  • ymodel.eval()El uso combinado produce los mejores resultados.

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

Otras extensiones