Fonction torch.inference_mode de PyTorch

Pytorch torch 参考手册Manuel de référence Pytorch torch

torch.inference_modeest un gestionnaire de contexte dans PyTorch utilisé pour le mode inférence. Il est plustorch.no_gradstrict, non seulement en désactivant le calcul du gradient, mais aussi en désactivant toutes les fonctions de suivi du moteur autograd.

Ceci est, lors de l’inférence du modèle, plusno_gradefficace, et peut en outre réduire l’utilisation de la mémoire et améliorer la vitesse d’inférence.

Définition de fonction

torch.inference_mode(mode=True)

Parameters:

  • mode(bool, facultatif) : SiTrue(valeur par défaut), active le mode d’inférence ; siFalse, quitte le mode d’inférence. Valeur par défaut :True。

Valeur de retour:

  • Renvoie un gestionnaire de contexte dans lequel les gradients et autograd sont désactivés.

Exemples d'utilisation

Exemple 1 : Utilisation de base

Exemple

import torch

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

# Dans le contexte inference_mode
with torch.inference_mode():
    y = x * 2
    print("In inference_mode:", y.requires_grad)

# Dans le contexte no_grad
with torch.no_grad():
    z = x * 2
    print("In no_grad:", z.requires_grad)

La sortie est :

在 inference_mode 中: False
在 no_grad 中: False

Exemple 2 : Comparaison de no_grad et inference_mode

Exemple

import torch

# Créer un tenseur
x = torch.randn(100, 100)

# Dans inference_mode
with torch.inference_mode():
    # Beaucoup de calculs effectués
    for _ in range(10):
        x = torch.mm(x, x)

    # Même après que le calcul est terminé, les tenseurs dans le contexte ne peuvent pas être utilisés pour la rétropropagation
    result = x.sum()

    # Vérifier s'il peut être converti en un tenseur qui nécessite des gradients
    print("En inference_mode :", result.is_leaf)

# Effectuer le même calcul dans no_grad
x2 = torch.randn(100, 100)
with torch.no_grad():
    for _ in range(10):
        x2 = torch.mm(x2, x2)
    result2 = x2.sum()
    print("En no_grad :", result2.is_leaf)

La sortie est :

在 inference_mode 中: False
在 no_grad 中: True

Exemple 3 : Inférence de modèle

Exemple

import torch
import torch.nn as nn

# Définir un modèle simple
model = nn.Sequential(
    nn.Linear(10, 20),
    nn.ReLU(),
    nn.Linear(20, 5)
)

model.eval()

# Créer les données d'entrée
x = torch.randn(100, 10)

# Utiliser inference_mode pour l'inférence
with torch.inference_mode():
    output = model(x)
    print("Forme de sortie :", output.shape)
    print("requires_grad de la sortie :", output.requires_grad)

# On peut aussi utiliser le décorateur
@torch.inference_mode()
def predict(x):
    return model(x)

result = predict(x)
print("Méthode décorateur - Forme de sortie :", result.shape)

La sortie est :

输出形状: torch.Size([100, 5])
输出 requires_grad: False
装饰器方式 - 输出形状: torch.Size([100, 5])

Exemple 4 : Comparaison de l'optimisation de la mémoire

Exemple

import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Linear(1000, 1000),
    nn.ReLU(),
    nn.Linear(1000, 1000),
    nn.ReLU(),
    nn.Linear(1000, 10)
)

# Test de l'utilisation de la mémoire des différents modes
x = torch.randn(50, 1000)

print("Sans aucun gestionnaire de contexte :")
_ = model(x)

print("\nUtilisation de no_grad :")
with torch.no_grad():
    _ = model(x)

print("\nUtilisation de inference_mode :")
with torch.inference_mode():
    _ = model(x)

Utilisationinference_modepeut optimiser davantage la mémoire car il désactive complètement le moteur autograd.

Fonctions associées

  • torch.no_grad(): Désactive le calcul du gradient mais conserve certaines fonctionnalités d'autograd.
  • torch.enable_grad(): Active le calcul du gradient.
  • torch.is_inference_mode_enabled(): Vérifie si le mode d'inférence est activé.

Notes

  • inference_modeComparerno_gradplus strict, désactive davantage de fonctionnalités.
  • Dansinference_modeLes tenseurs qui y sont créés sont marqués comme nœuds non-feuilles et ne peuvent pas être utilisés pour la rétropropagation.
  • Il est recommandé de l'utiliser lors de l'inférence et de l'évaluation du modèleinference_modepour obtenir les meilleures performances.
  • inference_modeNe peut pas êtreno_gradimbriqué.

Pytorch torch 参考手册Manuel de référence Pytorch torch

Autres extensions