Fonction torch.stack de PyTorch

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

torch.stackC'est une fonction de PyTorch utilisée pour empiler plusieurs tenseurs le long d'une nouvelle dimension. Elle crée une nouvelle dimension et place tous les tenseurs d'entrée le long de cette nouvelle dimension.

Elle est couramment utilisée en apprentissage profond pour des scénarios tels que la création de lots de données et l'empilement des sorties de plusieurs modèles.

Définition de fonction

torch.stack(tensors, dim=0, out=None)

Parameters:

  • tensors(Séquence de Tensor) : La séquence de tenseurs à empiler. Tous les tenseurs doivent avoir la même forme.
  • dim(int, optionnel) : La dimension le long de laquelle empiler, par défaut 0. La nouvelle dimension sera insérée à cette position.
  • out(Tensor, optionnel) : Le tenseur de sortie.

Valeur de retour:

  • torch.Tensor: Renvoie le tenseur empilé.

Exemples d'utilisation

Exemple 1: Empilement de base

Exemple

import torch

# Créer deux tenseurs unidimensionnels
a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])

# Empiler selon une nouvelle dimension
c = torch.stack([a, b])

print("Forme de a :", a.shape)
print("Forme de b :", b.shape)
print("Forme de c :", c.shape)
print(c)

Le résultat de sortie est:

a 的形状: torch.Size([3])
b 的形状: torch.Size([3])
c 的形状: torch.Size([2, 3])
tensor([[1, 2, 3],
        [4, 5, 6]])

Exemple 2: Spécification de la dimension d'empilement

Exemple

import torch

a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])

# Empiler selon dim=1
c = torch.stack([a, b], dim=1)

print("Forme de c :", c.shape)
print(c)

Le résultat de sortie est:

c 的形状: torch.Size([3, 2])
tensor([[1, 4],
        [2, 5],
        [3, 6]])

Exemple 3: Empilement de plusieurs tenseurs

Exemple

import torch

# Créer plusieurs tenseurs
tensors = [torch.tensor([i, i+1, i+2]) for i in range(5)]

# Empiler tous les tenseurs
result = torch.stack(tensors)

print("Forme du résultat :", result.shape)
print(result)

Le résultat de sortie est:

结果形状: torch.Size([5, 3])
tensor([[0, 1, 2],
        [1, 2, 3],
        [2, 3, 4],
        [3, 4, 5],
        [4, 5, 6]])

Exemple 4: Sauvegarde de plusieurs états dans un réseau de neurones

Exemple

import torch

# Simuler l'enregistrement des valeurs de perte pour plusieurs époques
losses = []
for epoch in range(5):
    loss = torch.tensor([epoch * 0.1, (epoch + 1) * 0.1])
    losses.append(loss)

# Empiler les pertes de toutes les époques
all_losses = torch.stack(losses)

print("Forme de la perte au fil des époques :", all_losses.shape)
print(all_losses)

Le résultat de sortie est:

各 epoch 损失形状: torch.Size([5, 2])
tensor([[0.0000, 0.1000],
        [0.1000, 0.2000],
        [0.2000, 0.3000],
        [0.3000, 0.4000],
        [0.4000, 0.5000]])

Différence entre torch.stack et torch.cat

Exemple

import torch

a = torch.randn(2, 3)
b = torch.randn(2, 3)

# stack crée une nouvelle dimension
stack_result = torch.stack([a, b])
print("forme du résultat de stack :", stack_result.shape)

# cat ne crée pas de nouvelle dimension
cat_result = torch.cat([a, b], dim=0)
print("forme du résultat de cat :", cat_result.shape)

Le résultat de sortie est:

stack 结果形状: torch.Size([2, 2, 3])
cat 结果形状: torch.Size([4, 3])
  • torch.stack: Empile le long d'une nouvelle dimension ; les tenseurs d'entrée doivent avoir exactement la même forme, et une nouvelle dimension est ajoutée.
  • torch.cat: Concatène le long d'une dimension existante ; les tenseurs d'entrée peuvent différer dans la dimension de concaténation, et aucune nouvelle dimension n'est ajoutée.

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

Autres extensions