LSTM PyTorch : dimensions, états et sorties de séquence

La couche nn.LSTM de PyTorch traite une séquence et renvoie deux éléments : une sortie à chaque pas de temps issue de la dernière couche, et une paire d’états finaux, (h_n, c_n). Avec batch_first=True, l’entrée par lots suit l’ordre (lot, temps, caractéristiques). Les états finaux conservent une autre disposition, avec les couches et les directions en premier.

Ce guide utilise cinq exemples CPU exécutables pour distinguer ces tenseurs. Nous avons effectué 37 contrôles avec Python 3.12.14 et PyTorch 2.15.0+cpu, dont un calcul scalaire indépendant, des séquences de longueurs variables et la transmission d’état entre segments. Chaque bloc Python s’exécute indépendamment avec PyTorch installé. Ces exemples vérifient le fonctionnement, sans démontrer une précision de prévision ni des performances GPU.

Commencer avec une seule unité cachée

Un LSTM conserve un état caché et un état de cellule. Dans ce premier exemple, nous choisissons les coefficients pour pouvoir vérifier le calcul à la main. Tous les poids et biais sont nuls, sauf le poids d’entrée de la valeur candidate de la cellule. La référence de l’API LSTM indique l’ordre des portes utilisé pour trouver ce coefficient.

import torch

x = torch.tensor([[[1.0], [2.0], [3.0]]])
layer = torch.nn.LSTM(1, 1, batch_first=True)
with torch.no_grad():
    for parameter in layer.parameters():
        parameter.zero_()
    # Gate rows are input, forget, candidate, output.
    layer.weight_ih_l0[2, 0] = 1.0
output, (h_n, c_n) = layer(x)
print('Output shape:', tuple(output.shape))
print('Hidden values:', [round(v, 6) for v in output.flatten().tolist()])
print('Final hidden:', round(h_n.item(), 6))
print('Final cell:', round(c_n.item(), 6))

Sortie observée :

Output shape: (1, 3, 1)
Hidden values: [0.1817, 0.293282, 0.341238]
Final hidden: 0.341238
Final cell: 0.833734

Avec ces coefficients, les portes d’entrée, d’oubli et de sortie valent chacune 0,5. En partant d’une cellule nulle, notre calcul indépendant est c_t = 0.5 × c_(t-1) + 0.5 × tanh(x_t), puis h_t = 0.5 × tanh(c_t). Pour la première entrée, 1, la cellule vaut environ 0,380797 et l’état caché 0,181700.

Après les trois entrées, la cellule vaut 0,833734 et l’état caché 0,341238. Ce sont des quantités différentes. Le dernier élément de cette sortie unidirectionnelle correspond à l’état caché final, pas à la cellule finale. Ces poids rendent la récurrence visible ; ils ne constituent pas un modèle de séquence entraîné.

Distinguer les axes du lot et ceux des états

Supposons que chaque exemple comporte quatre pas de temps et trois mesures par pas. Un lot de deux exemples a la forme (2, 4, 3). Ici, input_size=3 désigne les caractéristiques d’un pas, pas la taille du lot ni la longueur de la séquence.

Nous créons deux couches aux paramètres identiques et ne changeons que la disposition de leur entrée. Transposer les axes du lot et du temps permet de comparer les résultats :

import torch

torch.manual_seed(7)
x = torch.arange(24, dtype=torch.float32).reshape(2, 4, 3) / 10
batch_layer = torch.nn.LSTM(3, 5, num_layers=2, batch_first=True)
time_layer = torch.nn.LSTM(3, 5, num_layers=2)
time_layer.load_state_dict(batch_layer.state_dict())
batch_out, (batch_h, batch_c) = batch_layer(x)
time_out, (time_h, time_c) = time_layer(x.transpose(0, 1))
print('Batch-first output:', tuple(batch_out.shape))
print('Time-first output:', tuple(time_out.shape))
print('Hidden and cell:', tuple(batch_h.shape), tuple(batch_c.shape))
print('Same output:', torch.allclose(batch_out, time_out.transpose(0, 1)))
print('Same final states:', torch.allclose(batch_h, time_h)
      and torch.allclose(batch_c, time_c))

Sortie observée :

Batch-first output: (2, 4, 5)
Time-first output: (4, 2, 5)
Hidden and cell: (2, 2, 5) (2, 2, 5)
Same output: True
Same final states: True

La sortie avec le lot en premier a la forme (2, 4, 5) : deux séquences, quatre positions et cinq caractéristiques cachées. Celle avec le temps en premier a la forme (4, 2, 5). Après transposition, toutes les valeurs correspondent dans la tolérance de comparaison.

Les deux états finaux ont ici la forme (2, 2, 5) : deux couches, deux éléments dans le lot et une largeur cachée de cinq. Modifier batch_first ne déplace pas l’axe du lot dans ces états. Des contrôles séparés ont confirmé que des états initiaux omis correspondent à des tenseurs explicitement nuls, et qu’un nombre incorrect de couches dans l’état initial provoque une erreur.

Utilisez une transposition des axes plutôt qu’un simple reshape pour passer d’un ordre à l’autre. Notre contre-exemple avec reshape a les dimensions voulues, mais place les valeurs différemment. Vérifiez les lots produits par votre DataLoader, dont notre guide est en anglais, avant d’ajouter une conversion.

Lire les deux directions dans les états finaux

L’exemple suivant utilise deux couches, deux directions et une largeur cachée de trois. La largeur de sortie vaut six, car chaque position contient les caractéristiques des deux directions. Le premier axe des états comporte quatre entrées : deux directions pour chacune des deux couches.

import torch

torch.manual_seed(7)
x = torch.arange(16, dtype=torch.float32).reshape(2, 4, 2) / 10
layer = torch.nn.LSTM(2, 3, num_layers=2, bidirectional=True, batch_first=True)
output, (h_n, c_n) = layer(x)
# Separate layer, direction, batch, hidden width.
states = h_n.reshape(2, 2, 2, 3)
summary = torch.cat((states[-1, 0], states[-1, 1]), dim=-1)
print('Output:', tuple(output.shape))
print('Hidden and cell:', tuple(h_n.shape), tuple(c_n.shape))
print('Summary:', tuple(summary.shape))
print('Forward endpoint:', torch.allclose(output[:, -1, :3], states[-1, 0]))
print('Reverse endpoint:', torch.allclose(output[:, 0, 3:], states[-1, 1]))
print('Last output is full summary:', torch.allclose(output[:, -1], summary))

Sortie observée :

Output: (2, 4, 6)
Hidden and cell: (4, 2, 3) (4, 2, 3)
Summary: (2, 6)
Forward endpoint: True
Reverse endpoint: True
Last output is full summary: False

En séparant les axes de couche et de direction dans h_n, nous pouvons sélectionner la dernière couche et concaténer ses deux états finaux. Le résultat a la forme (2, 6), soit six valeurs par élément du lot. Ce reshape divise un axe combiné ; il ne réorganise pas les axes du lot et du temps comme dans l’avertissement précédent.

L’état final dans le sens direct correspond à la première moitié de la sortie à la dernière position. L’état final dans le sens inverse correspond à la seconde moitié à la première position. La dernière position de sortie, seule, échoue à notre comparaison avec le résumé complet. Le raccourci output[:, -1] est donc incorrect pour cet exemple bidirectionnel.

Un modèle qui utilise des éléments futurs de la séquence ne peut pas fournir le même résultat bidirectionnel avant leur arrivée. Choisissez les directions selon les informations disponibles au moment de la prédiction. Pour comparer les modèles récurrents à d’autres approches, consultez notre guide des architectures de réseaux neuronaux, en anglais.

Exclure le remplissage du calcul récurrent

Le lot suivant contient des séquences de deux et quatre pas. Nous complétons la plus courte avec des zéros pour obtenir un tenseur rectangulaire. Le transmettre directement au LSTM effectue quatre pas pour les deux exemples : une entrée nulle reste un pas de récurrence.

pack_padded_sequence permet d’indiquer les longueurs valides. Le tenseur des longueurs reste sur CPU. Nous utilisons enforce_sorted=False parce que les longueurs ne sont pas décroissantes. Chaque paramètre est fixé à 0,1 pour rendre la comparaison reproductible.

import torch

from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

x = torch.tensor([[[1.], [2.], [0.], [0.]],
                  [[3.], [4.], [5.], [6.]]])
lengths = torch.tensor([2, 4], dtype=torch.long)  # CPU tensor
layer = torch.nn.LSTM(1, 2, batch_first=True)
with torch.no_grad():
    for parameter in layer.parameters():
        parameter.fill_(0.1)
packed = pack_padded_sequence(x, lengths, batch_first=True, enforce_sorted=False)
packed_out, (packed_h, packed_c) = layer(packed)
padded_out, recovered = pad_packed_sequence(packed_out, batch_first=True, total_length=4)
_, (single_h, single_c) = layer(x[:1, :2])
_, (plain_h, plain_c) = layer(x)
print('Recovered lengths:', recovered.tolist())
print('Padded output:', tuple(padded_out.shape))
print('Short sequence matches alone:', torch.allclose(packed_h[:, :1], single_h))
print('Plain padding matches:', torch.allclose(plain_h[:, :1], single_h))
print('Unpacked tail is zero:', torch.count_nonzero(padded_out[0, 2:]).item() == 0)

Sortie observée :

Recovered lengths: [2, 4]
Padded output: (2, 4, 2)
Short sequence matches alone: True
Plain padding matches: False
Unpacked tail is zero: True

L’état caché de la séquence courte empaquetée correspond au traitement de ses deux vrais pas seuls. Un contrôle supplémentaire a confirmé la même correspondance pour l’état de cellule. Le remplissage ordinaire avec des zéros ne donne pas cet état caché final : les deux pas ajoutés l’ont modifié.

Nous avons aussi vérifié séparément la séquence longue et comparé chaque valeur de sortie valide pour les deux exemples. Les états renvoyés conservent l’ordre initial du lot malgré le tri interne effectué lors de l’empaquetage. C’est nécessaire pour associer les représentations aux bonnes étiquettes ou aux bons identifiants.

pad_packed_sequence rétablit une sortie rectangulaire et renvoie les longueurs. Dans notre résultat, la partie inutilisée de la séquence courte est nulle, et total_length=4 produit la dimension temporelle attendue. Ces zéros sont du remplissage de sortie, pas deux calculs récurrents supplémentaires.

Les longueurs désignent les pas valides au début de chaque séquence. Un masque comportant des trous au milieu ne peut pas être traité comme une simple longueur. Nos tests ont rejeté une longueur nulle ainsi que des longueurs non triées avec l’exigence de tri par défaut. Empaqueter les entrées n’exclut pas automatiquement les positions remplies d’une perte calculée après dépaquetage ; sélectionnez explicitement les positions valides.

Transmettre l’état entre segments d’une même séquence

Pour un LSTM unidirectionnel dont les poids restent inchangés, sans dropout, nous divisons une séquence en deux segments. Transmettre les deux états du premier segment au second reproduit le résultat de la séquence entière dans ce test :

import torch

torch.manual_seed(7)
x = torch.arange(12, dtype=torch.float32).reshape(1, 6, 2) / 10
layer = torch.nn.LSTM(2, 3, batch_first=True)
whole, whole_state = layer(x)
first, state = layer(x[:, :3])
second, final_state = layer(x[:, 3:], state)
reset_second, _ = layer(x[:, 3:])
print('Carried output matches:', torch.allclose(torch.cat((first, second), dim=1), whole))
print('Carried states match:', all(torch.allclose(a, b) for a, b in zip(final_state, whole_state)))
print('Reset output matches:', torch.allclose(reset_second, whole[:, 3:]))

Sortie observée :

Carried output matches: True
Carried states match: True
Reset output matches: False

Réinitialiser l’état au début du second segment modifie sa sortie. Le transmettre conserve le lien entre les segments. Chaque position du lot doit toujours désigner la même séquence : ne transmettez pas aveuglément l’état entre des lots mélangés ou des exemples indépendants.

Pour tronquer l’historique de calcul pendant l’entraînement, tuple(s.detach() for s in state) conserve les valeurs des états tout en détachant leur graphe antérieur. Nous avons vérifié ces deux propriétés localement. Cette opération change les pas antérieurs qui peuvent recevoir des gradients ; elle ne remet pas les valeurs à zéro et ne démontre pas la convergence d’une stratégie d’entraînement.

Cette comparaison par segments ne couvre ni le traitement bidirectionnel, ni le dropout, ni les mises à jour de paramètres entre segments. Pour reprendre un entraînement, distinguez ces états de séquence d’un point de contrôle complet du modèle et de l’optimiseur. Notre guide de sauvegarde et de chargement, en anglais traite cette reprise.

Vérifier ces points avant l’entraînement

• Largeur des caractéristiques : la dernière dimension de l’entrée doit correspondre à input_size. Notre test de largeur incorrecte a provoqué une erreur.
• Axes des états : identifiez couches, directions, lot et largeur cachée avant d’indexer h_n.
• Limites des séquences : utilisez les vraies longueurs ou des réinitialisations choisies ; remplissage et exemples indépendants ne sont pas des continuations.
• Sens de la sortie : une représentation LSTM n’est pas encore une probabilité de classe ni une prévision. Définissez séparément la tête de prédiction et la cible.

Tous les exemples désactivent les projections et conservent un dropout nul. Ne recopiez pas leurs largeurs d’état dans un LSTM avec projection sans vérifier sa configuration. Nous avons aussi comparé une entrée sans axe de lot à un lot d’un élément et vérifié des gradients finis pour l’entrée et les paramètres lors d’une rétropropagation.

Les tests établissent les valeurs affichées, les dispositions et le comportement des états. Ils n’évaluent aucun jeu de données, précision métier, CUDA, précision mixte, débit ni stabilité d’entraînement à long terme. L’environnement a signalé l’absence de NumPy à l’initialisation ; ces contrôles reposant uniquement sur des tenseurs se sont terminés sans NumPy.