PyTorch : torch.cat et torch.stack, exemples et dimensions

Utilisez torch.cat pour joindre des tenseurs selon une dimension existante. Utilisez torch.stack pour insérer une nouvelle dimension qui sépare les entrées. Pour deux matrices de dimensions (2, 3), cat peut produire (4, 3) ou (2, 6) ; stack produit un tenseur à trois dimensions.

La question principale est le sens de chaque axe : échantillons, pas de temps, canaux, caractéristiques ou autre. Une forme de sortie valide ne prouve pas que les valeurs sont disposées comme l’attend la couche suivante.

Ce guide suit cinq exemples originaux avec leurs sorties capturées. Nous avons exécuté 36 vérifications avec Python 3.12.14 et PyTorch 2.15.0+cpu. Les données entières rendent la disposition visible ; les données float64 vérifient les gradients. Il s’agit de tests de comportement sur CPU, pas de mesures de performances.

Joindre des lignes ou des colonnes avec torch.cat

Les deux entrées ci-dessous contiennent deux lignes et trois colonnes. Les valeurs de la seconde matrice sont plus grandes pour distinguer facilement leur provenance.

import torch

a = torch.tensor([[1, 2, 3], [4, 5, 6]])
b = torch.tensor([[10, 20, 30], [40, 50, 60]])
rows = torch.cat((a, b), dim=0)
columns = torch.cat((a, b), dim=1)
last_axis = torch.cat((a, b), dim=-1)
print("dim=0:", tuple(rows.shape), rows.tolist())
print("dim=1:", tuple(columns.shape), columns.tolist())
print("dim=-1 equals dim=1:", torch.equal(last_axis, columns))

dim=0: (4, 3) [[1, 2, 3], [4, 5, 6], [10, 20, 30], [40, 50, 60]]
dim=1: (2, 6) [[1, 2, 3, 10, 20, 30], [4, 5, 6, 40, 50, 60]]
dim=-1 equals dim=1: True

Avec dim=0, les nombres de lignes s’additionnent : deux plus deux donne quatre, tandis que le nombre de colonnes reste trois. Avec dim=1, les colonnes s’additionnent et produisent deux lignes de six valeurs. L’opération conserve deux dimensions dans les deux cas.

Pour cette entrée à deux dimensions, dim=-1 désigne le dernier axe et correspond à dim=1. Dans un tenseur plus grand, il ne désigne pas toujours des colonnes : il désigne l’axe placé en dernier.

La référence de torch.cat impose des dimensions identiques hors de l’axe de concaténation, avec une exception documentée pour un tenseur vide à une dimension de forme (0,). Les exemples suivants montrent les conséquences pratiques de cette règle.

Créer un nouvel axe avec torch.stack

L’empilement conserve une position distincte pour chaque entrée sur l’axe ajouté. Toutes les entrées doivent avoir la même forme. Dans l’API torch.stack, dim choisit la position de ce nouvel axe.

import torch

a = torch.tensor([[1, 2, 3], [4, 5, 6]])
b = torch.tensor([[10, 20, 30], [40, 50, 60]])
front = torch.stack((a, b), dim=0)
middle = torch.stack((a, b), dim=1)
back = torch.stack((a, b), dim=-1)
via_cat = torch.cat((a.unsqueeze(1), b.unsqueeze(1)), dim=1)
print("dim=0:", tuple(front.shape), front.tolist())
print("dim=1:", tuple(middle.shape), middle.tolist())
print("dim=-1:", tuple(back.shape), back.tolist())
print("unsqueeze then cat equals stack:", torch.equal(via_cat, middle))

dim=0: (2, 2, 3) [[[1, 2, 3], [4, 5, 6]], [[10, 20, 30], [40, 50, 60]]]
dim=1: (2, 2, 3) [[[1, 2, 3], [10, 20, 30]], [[4, 5, 6], [40, 50, 60]]]
dim=-1: (2, 3, 2) [[[1, 10], [2, 20], [3, 30]], [[4, 40], [5, 50], [6, 60]]]
unsqueeze then cat equals stack: True

Avec dim=0, la première tranche de sortie contient toute la première matrice, et la seconde toute la deuxième. Avec dim=1, chaque ligne d’origine contient une paire de lignes provenant des deux sources. Les deux sorties ont ici la même forme (2, 2, 3), car le nombre d’entrées et le nombre initial de lignes valent tous deux deux. Leurs valeurs sont pourtant disposées différemment.

Avec dim=-1, la sortie a pour forme (2, 3, 2). Chaque paire finale contient les valeurs correspondantes des deux entrées, par exemple [1, 10]. Pour ces matrices, stack peut insérer un axe aux positions positives 0, 1 ou 2. Cat ne dispose que des axes existants 0 et 1.

La dernière ligne vérifie une équivalence utile : ajouter un axe de taille un à chaque entrée, puis concaténer sur cet axe, reproduit cet empilement. unsqueeze ajoute cet axe et partage les données sous-jacentes du tenseur source. L’ajout de l’axe ne réorganise pas les valeurs initiales.

Choisir l’axe selon le sens des données

Si deux tenseurs contiennent des lots de vecteurs de caractéristiques de forme (batch, features), cat avec dim=0 réunit les lots en conservant la largeur des caractéristiques. Avec dim=1, il joint les caractéristiques des lignes correspondantes. Ce second usage suppose que les lignes décrivent les mêmes échantillons, dans le même ordre.

Les dimensions seules ne prouvent pas cette correspondance. Deux lots peuvent avoir la même longueur et contenir des échantillons sans rapport. La jointure de leurs colonnes fonctionnerait, mais les enregistrements obtenus auraient un sens incorrect.

Si chaque entrée représente un échantillon de même forme, stack avec dim=0 peut créer un axe de lot. Si les entrées sont déjà des lots, ce même appel ajoute un axe autour de ces lots. Inspectez les tenseurs renvoyés par votre DataLoader PyTorch, présenté dans notre guide en anglais, avant de choisir.

Traiter les longueurs différentes sans supposer de broadcasting

Un dernier lot plus court peut encore être concaténé si sa largeur de caractéristiques correspond :

import torch

a = torch.tensor([[1, 2, 3], [4, 5, 6]])
short = torch.tensor([[7, 8, 9]])
joined = torch.cat((a, short), dim=0)
print("different batch sizes:", tuple(joined.shape), joined.tolist())
errors = {}
for label, operation in [
    ("stack unequal shapes", lambda: torch.stack((a, short))),
    ("cat unequal non-join axes", lambda: torch.cat((a, short), dim=1)),
    ("cat missing axis", lambda: torch.cat((a, a), dim=2)),
]:
    try:
        operation()
    except (RuntimeError, IndexError) as exc:
        errors[label] = type(exc).__name__
        print(label + ":", type(exc).__name__)

different batch sizes: (3, 3) [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
stack unequal shapes: RuntimeError
cat unequal non-join axes: RuntimeError
cat missing axis: IndexError

Le premier appel réunit un tenseur (2, 3) et un tenseur (1, 3) en un résultat (3, 3). Leur empilement échoue parce que leurs formes diffèrent. La concaténation avec dim=1 échoue aussi : l’axe des lignes, qui n’est pas concaténé, a des longueurs de deux et de un.

Cat ne répète pas automatiquement l’entrée d’une ligne pour l’adapter à l’autre tenseur. L’erreur indique que l’alignement souhaité doit être examiné. Ajouter du padding, utiliser un format de séquence compacté ou répéter explicitement des valeurs sont des choix distincts, avec des conséquences sur les données.

Le dernier appel tente de concaténer sur dim=2, alors que les entrées ne possèdent que deux axes. Il provoque IndexError. Le texte des erreurs peut varier selon la version ; cet exemple affiche la classe de l’exception dans l’environnement testé.

Nos vérifications supplémentaires rejettent une liste vide passée à l’une ou l’autre fonction. Un tenseur sans lignes constitue un autre cas : concaténer un tenseur (0, 3) avec notre entrée (2, 3), sur dim=0, n’ajoute aucune ligne. Nous avons aussi vérifié l’exception documentée du tenseur vide (0,) pour cat. Distinguez une collection d’entrées vide d’un tenseur d’entrée sans éléments.

Empiler des scalaires ou ajouter un axe avant la concaténation

Un tenseur scalaire a pour forme () : il ne possède aucun axe existant à prolonger. L’exemple suivant rend cette différence explicite :

import torch

x = torch.tensor(3.0)
y = torch.tensor(7.0)
stacked = torch.stack((x, y))
promoted = torch.cat((x.unsqueeze(0), y.unsqueeze(0)))
print("scalar input shape:", tuple(x.shape))
print("stacked:", tuple(stacked.shape), stacked.tolist())
print("promoted cat:", tuple(promoted.shape), promoted.tolist())
try:
    torch.cat((x, y))
except RuntimeError as exc:
    scalar_error = type(exc).__name__
    print("cat scalars:", scalar_error)

scalar input shape: ()
stacked: (2,) [3.0, 7.0]
promoted cat: (2,) [3.0, 7.0]
cat scalars: RuntimeError

Stack crée un axe de vecteur et renvoie [3.0, 7.0]. Cat échoue lorsqu’il reçoit directement les scalaires. Ajouter d’abord un axe de taille un à chacun produit deux vecteurs d’un élément, que cat peut réunir dans le même résultat.

Ce cas est utile pour regrouper des résultats scalaires qui sont encore des tenseurs. Déterminez si leurs connexions de gradient doivent être conservées avant de les convertir en nombres Python ordinaires ; l’exemple suivant conserve des tenseurs tout au long du calcul.

Vérifier le retour des gradients vers les entrées

Le résultat joint peut participer à un calcul autograd. Nous utilisons deux pertes différentes pour vérifier indépendamment les gradients attendus aux entrées.

import torch

x = torch.tensor([1.0, 2.0], dtype=torch.float64, requires_grad=True)
y = torch.tensor([3.0, 4.0], dtype=torch.float64, requires_grad=True)
torch.cat((x, y)).square().sum().backward()
cat_x_grad = x.grad.clone()
cat_y_grad = y.grad.clone()
print("cat gradients:", cat_x_grad.tolist(), cat_y_grad.tolist())

x.grad = None
y.grad = None
weights = torch.tensor([[1.0, 10.0], [100.0, 1000.0]], dtype=torch.float64)
(torch.stack((x, y), dim=1) * weights).sum().backward()
stack_x_grad = x.grad.clone()
stack_y_grad = y.grad.clone()
print("stack gradients:", stack_x_grad.tolist(), stack_y_grad.tolist())

x.grad = None
torch.cat((x, x)).sum().backward()
print("same input used twice:", x.grad.tolist())

cat gradients: [2.0, 4.0] [6.0, 8.0]
stack gradients: [1.0, 100.0] [10.0, 1000.0]
same input used twice: [2.0, 2.0]

Pour cat, la perte est la somme des carrés des valeurs jointes. Sa dérivée vaut deux fois chaque valeur : x reçoit donc [2, 4] et y reçoit [6, 8]. Pour stack avec dim=1, x occupe la première colonne et y la seconde. La multiplication par la matrice de coefficients donne les gradients [1, 100] pour x et [10, 1000] pour y.

Nous effaçons les anciens gradients entre les calculs. Dans le dernier, x apparaît deux fois dans la concaténation : chaque élément reçoit donc deux contributions de un. Le gradient obtenu est [2, 2]. Les tests confirment aussi que backward laisse les valeurs sources inchangées.

Joindre des tenseurs est distinct de l’accumulation des gradients sur plusieurs micro-lots. Notre guide de l’accumulation des gradients, en anglais, traite de la pondération des pertes et du moment des étapes d’optimiseur. Réunir les valeurs des tenseurs ne suffit pas à mettre en œuvre cette boucle d’entraînement.

Vérifier une valeur connue en plus des dimensions

Avant de joindre des données réelles, notez le sens de chaque axe et la forme attendue en sortie. Suivez ensuite une valeur connue de chaque entrée jusqu’à sa destination. Les deux sorties stack de forme (2, 2, 3) montrent pourquoi le contrôle des dimensions seul peut manquer une erreur d’axe.

Ces tests couvrent des tenseurs denses sur CPU, certains cas invalides et le retour exact des gradients. Une entrée non contiguë fonctionne aussi dans un test cat distinct. Nous n’avons évalué ni les performances CUDA, ni les tenseurs creux ou quantifiés, ni les entrées sur plusieurs périphériques, ni autograd avec l’argument out. Les performances et la qualité d’entraînement demandent des vérifications distinctes.