Fonction torch.enable_grad de PyTorch

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

torch.enable_gradC'est un gestionnaire de contexte dans PyTorch utilisé pour activer le calcul du gradient. Il esttorch.no_gradl'opposé, et est utilisé pour activer explicitement le calcul du gradient dans les blocs de code qui nécessitent des gradients.

C'est très utile lorsque vous devez temporairement entraîner un modèle dans du code d'inférence, ou lorsque vous devez activer le calcul du gradient dans une région locale.

Définition de fonction

torch.enable_grad()

Parameters:

  • Aucun paramètre. C'est un gestionnaire de contexte.

Valeur de retour:

  • Renvoie un gestionnaire de contexte qui active le calcul du gradient dans ce contexte.

Exemples d'utilisation

Exemple 1 : Utilisation de base

Exemple

import torch

# Par défaut, dans un contexte no_grad
with torch.no_grad():
    x = torch.tensor([1.0, 2.0, 3.0])
    print("In no_grad:", x.requires_grad)

    # Utiliser enable_grad pour activer temporairement les gradients
    with torch.enable_grad():
        y = x * 2
        print("Dans enable_grad :", y.requires_grad)

    # Après la sortie, restaurer l'état no_grad
    z = x * 2
    print("Après la sortie :", z.requires_grad)

La sortie est :

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

Exemple 2 : Mélange d'entraînement et d'inférence

Exemple

import torch
import torch.nn as nn

model = nn.Linear(10, 2)

# Mode inférence
with torch.no_grad():
    x = torch.randn(5, 10)
    output1 = model(x)
    print("Sortie d'inférence :", output1.shape)

    # Si vous devez entraîner temporairement certains paramètres pendant l'inférence
    with torch.enable_grad():
        # Créer un tenseur qui nécessite des gradients pour le calcul
        temp_weight = torch.randn(10, 10, requires_grad=True)
        temp_output = torch.mm(x, temp_weight)
        print("Activer temporairement les gradients :", temp_output.requires_grad)

La sortie est :

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

Exemple 3 : Utilisation comme décorateur

Exemple

import torch

@torch.enable_grad()
def train_step(x, y):
    """Simuler les étapes d'entraînement"""
    # Cette fonction activera le calcul des gradients en interne
    loss = (x - y).sum()
    return loss

# Appelez-la dans un contexte no_grad
with torch.no_grad():
    x = torch.tensor([1.0, 2.0, 3.0])
    y = torch.tensor([0.0, 0.0, 0.0])
    # Le décorateur garantit que le calcul des gradients est activé
    loss = train_step(x, y)
    print("Loss:", loss)
    print("Requires gradient:", loss.requires_grad)

La sortie est :

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

Fonctions associées

  • torch.no_grad(): Désactive le calcul des gradients.
  • torch.set_grad_enabled(grad): Active ou désactive le calcul des gradients selon les paramètres.
  • torch.is_grad_enabled(): Vérifie si le calcul du gradient est actuellement activé.

Notes

  • enable_gradPrincipalement utilisé pourno_gradactiver temporairement le calcul du gradient dans un contexte.
  • S'il est utilisé dans un environnement où les gradients sont globalement activésenable_grad, cela n'aura aucun effet.
  • Il est recommandé d'utiliserenable_gradle décorateur pour garantir que les gradients sont toujours activés à l'intérieur de la fonction.

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

Autres extensions