Función PyTorch torch.is_inference_mode_enabled

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

torch.is_inference_mode_enabledEs la función en PyTorch que se utiliza para comprobar si el modo de inferencia está habilitado actualmente. Devuelve el valor booleano de si se encuentra actualmente eninference_modeel contexto.

Esto es muy útil al escribir código que necesita ejecutar lógica diferente según el estado del modo de inferencia.

Definición de la función

torch.is_inference_mode_enabled()

Parámetros:

  • Sin parámetros.

Valor de retorno:

  • Devuelve un valor booleano: si el modo de inferencia está habilitado actualmente, devuelveTrue; de lo contrario, devuelveFalse。

Ejemplos de uso

Ejemplo 1: Uso básico

Ejemplo

import torch

# Por defecto, el modo de inferencia está deshabilitado
print("Estado por defecto:", torch.is_inference_mode_enabled())

# Dentro del contexto inference_mode
with torch.inference_mode():
    print("En inference_mode:", torch.is_inference_mode_enabled())

# Restaurado después de salir
print("Después de salir:", torch.is_inference_mode_enabled())

El resultado es:

默认状态: False
在 inference_mode 中: True
退出后: False

Ejemplo 2: Comparación entre is_grad_enabled y is_inference_mode_enabled

Ejemplo

import torch

# En no_grad
with torch.no_grad():
    print("En no_grad:")
    print("  is_grad_enabled:", torch.is_grad_enabled())
    print("  is_inference_mode_enabled:", torch.is_inference_mode_enabled())

# En inference_mode
with torch.inference_mode():
    print("En inference_mode:")
    print("  is_grad_enabled:", torch.is_grad_enabled())
    print("  is_inference_mode_enabled:", torch.is_inference_mode_enabled())

# Estado por defecto
print("Estado por defecto:")
print("  is_grad_enabled:", torch.is_grad_enabled())
print("  is_inference_mode_enabled:", torch.is_inference_mode_enabled())

El resultado es:

在 no_grad 中:
  is_grad_enabled: False
  is_inference_mode_enabled: False
在 inference_mode 中:
  is_grad_enabled: False
  is_inference_mode_enabled: True
默认状态:
  is_grad_enabled: True
  is_inference_mode_enabled: False

Ejemplo 3: Uso en sentencias condicionales

Ejemplo

import torch
import torch.nn as nn

def forward_pass(x, model):
    """Ejecutar diferentes optimizaciones según el modo de inferencia"""
    if torch.is_inference_mode_enabled():
        print("Optimización usando modo de inferencia")
        # Cálculo optimizado bajo el modo de inferencia
        return model(x)
    elif torch.is_grad_enabled():
        print("Modo entrenamiento")
        return model(x)
    else:
        print("Modo eval")
        return model(x)

model = nn.Linear(10, 5)
x = torch.randn(5, 10)

# Probar diferentes estados
print("=== Modo entrenamiento ===")
with torch.enable_grad():
    result = forward_pass(x, model)

print("n=== Modo de inferencia ===")
with torch.inference_mode():
    result = forward_pass(x, model)

print("n=== Modo eval ===")
with torch.no_grad():
    result = forward_pass(x, model)

El resultado es:

=== 训练模式 ===
训练模式

=== 推理模式 ===
使用推理模式优化

=== eval 模式 ===
eval 模式

Ejemplo 4: Detección en un módulo personalizado

Ejemplo

import torch
import torch.nn as nn

class OptimizedLayer(nn.Module):
    def __init__(self):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(20, 10))

    def forward(self, x):
        # Detectar el modo de inferencia y aplicar optimizaciones
        if torch.is_inference_mode_enabled():
            # Modo de inferencia: utilizar un método de cálculo más eficiente
            return torch.mm(x, self.weight)
        elif torch.is_grad_enabled():
            # Modo de entrenamiento: conservar el cálculo de gradientes
            return torch.mm(x, self.weight)
        else:
            # Modo eval: no necesita gradientes pero puede usar optimizaciones
            with torch.no_grad():
                return torch.mm(x, self.weight)

layer = OptimizedLayer()
x = torch.randn(5, 10)

print("=== Modo de entrenamiento ===")
with torch.enable_grad():
    _ = layer(x)

print("n=== Modo de inferencia ===")
with torch.inference_mode():
    _ = layer(x)

print("n=== Modo eval ===")
with torch.no_grad():
    _ = layer(x)

El resultado es:

=== 训练模式 ===
训练模式

=== 推理模式 ===
推理模式：使用更高效的计算方式

=== eval 模式 ===
eval 模式

Funciones relacionadas

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

Notas

  • is_inference_mode_enabledEs una función de solo lectura y no cambia ningún estado.
  • Está diseñada específicamente para detectarinference_mode, en lugar deno_grad。
  • Eninference_modecuando se encuentra enis_grad_enabledtambién devuelveFalse, pero lo contrario no se cumple.
  • Al escribir código genérico, se puede usar esta función para distinguir diferentes modos de ejecución.

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

Otras extensiones