Fonction torch.set_grad_enabled de PyTorch

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

torch.set_grad_enabledest une fonction de PyTorch utilisée pour définir globalement si le calcul des gradients est activé. Elle peut activer ou désactiver dynamiquement le calcul des gradients ; contrairement aux gestionnaires de contexte, c'est une fonction qui peut modifier l'état global.

C'est très utile lorsque vous devez contrôler dynamiquement le calcul des gradients en fonction de conditions.

Définition de fonction

torch.set_grad_enabled(mode)

Parameters:

  • mode(bool) : S'il estTrue, active le calcul des gradients ; s'il estFalse, désactive le calcul des gradients.

Valeur de retour:

  • Renvoie un gestionnaire de contexte qui peut être utilisé danswithinstructions.

Exemples d'utilisation

Exemple 1 : Utilisation de base

Exemple

import torch

# Créer un tenseur qui nécessite des gradients
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)

# Désactiver le calcul du gradient
torch.set_grad_enabled(False)
y1 = x * 2
print("Après désactivation des gradients :", y1.requires_grad)

# Activer le calcul des gradients
torch.set_grad_enabled(True)
y2 = x * 2
print("Après activation des gradients :", y2.requires_grad)

La sortie est :

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

Exemple 2 : Utilisation comme gestionnaire de contexte

Exemple

import torch

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

# Utiliser un gestionnaire de contexte
with torch.set_grad_enabled(False):
    y1 = x * 2
    print("À l'intérieur du gestionnaire de contexte :", y1.requires_grad)

y2 = x * 2
print("En dehors du gestionnaire de contexte :", y2.requires_grad)

La sortie est :

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

Exemple 3 : Contrôle dynamique de l'entraînement et de l'inférence

Exemple

import torch
import torch.nn as nn

model = nn.Linear(10, 2)

def forward_pass(x, training=True):
    """Contrôler les gradients en fonction du paramètre d'entraînement"""
    with torch.set_grad_enabled(training):
        output = model(x)
        print(f"training={training}, requires_grad={output.requires_grad}")
    return output

# Mode entraînement
x = torch.randn(5, 10)
forward_pass(x, training=True)

# Mode inférence
forward_pass(x, training=False)

La sortie est :

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

Exemple 4 : Enregistrement et restauration de l'état des gradients

Exemple

import torch

# État initial
print("État initial :", torch.is_grad_enabled())

# Créer un gestionnaire de contexte qui revient à l'état d'origine
old = torch.is_grad_enabled()

# Désactiver temporairement les gradients
with torch.set_grad_enabled(False):
    print("À l'intérieur :", torch.is_grad_enabled())

# Restaurer automatiquement (mais ici nous devons restaurer manuellement)
torch.set_grad_enabled(old)
print("Après restauration :", torch.is_grad_enabled())

La sortie est :

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

Fonctions associées

  • torch.no_grad(): Gestionnaire de contexte qui désactive le calcul du gradient.
  • torch.enable_grad(): Gestionnaire de contexte qui active le calcul du gradient.
  • torch.is_grad_enabled(): Vérifie si le calcul du gradient est actuellement activé.

Notes

  • set_grad_enabledIl peut être appelé directement comme une fonction pour modifier l'état global, ou utilisé comme gestionnaire de contexte.
  • Lorsqu'il est appelé comme une fonction, vous devez restaurer manuellement l'état d'origine, sinon cela affectera le code suivant.
  • Il est recommandé d'utiliser l'approche du gestionnaire de contexte pour garantir que l'état est correctement restauré.
  • Soyez conscient de l'impact sur l'état global et utilisez-le avec prudence dans du code complexe.

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

Autres extensions