Función torch.stack de PyTorch

Pytorch torch 参考手册Manual de referencia de Pytorch torch

torch.stackEs una función en PyTorch que se utiliza para apilar múltiples tensores a lo largo de una nueva dimensión. Crea una nueva dimensión y coloca todos los tensores de entrada en esa nueva dimensión.

Esto se usa comúnmente en el aprendizaje profundo para crear datos por lotes, apilar las salidas de múltiples modelos, etc.

Definición de la función

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

Parámetros:

  • tensors(Sequence of Tensor): La secuencia de tensores a apilar. Todos los tensores deben tener la misma forma.
  • dim(int, opcional): La dimensión de apilamiento, el valor predeterminado es 0. La nueva dimensión se insertará en esta posición.
  • out(Tensor, opcional): El tensor de salida.

Valor de retorno:

  • torch.Tensor: Devuelve el tensor apilado.

Ejemplos de uso

Ejemplo 1: Apilamiento básico

Ejemplo

import torch

# Crear dos tensores unidimensionales
a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])

# Apilar a lo largo de la nueva dimensión
c = torch.stack([a, b])

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

El resultado es:

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

Ejemplo 2: Especificar la dimensión de apilamiento

Ejemplo

import torch

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

# Apilar a lo largo de dim=1
c = torch.stack([a, b], dim=1)

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

El resultado es:

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

Ejemplo 3: Apilar múltiples tensores

Ejemplo

import torch

# Crear múltiples tensores
tensors = [torch.tensor([i, i+1, i+2]) for i in range(5)]

# Apilar todos los tensores
result = torch.stack(tensors)

print("forma del resultado:", result.shape)
print(result)

El resultado es:

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

Ejemplo 4: Guardar múltiples estados en una red neuronal

Ejemplo

import torch

# Simular el guardado de los valores de pérdida de múltiples epochs
losses = []
for epoch in range(5):
    loss = torch.tensor([epoch * 0.1, (epoch + 1) * 0.1])
    losses.append(loss)

# Apilar las pérdidas de todos los epochs
all_losses = torch.stack(losses)

print("forma de la pérdida de cada epoch:", all_losses.shape)
print(all_losses)

El resultado es:

各 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]])

Diferencia entre torch.stack y torch.cat

Ejemplo

import torch

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

# stack crea una nueva dimensión
stack_result = torch.stack([a, b])
print("forma del resultado de stack:", stack_result.shape)

# cat no crea una nueva dimensión
cat_result = torch.cat([a, b], dim=0)
print("forma del resultado de cat:", cat_result.shape)

El resultado es:

stack 结果形状: torch.Size([2, 2, 3])
cat 结果形状: torch.Size([4, 3])
  • torch.stack: Apila a lo largo de una nueva dimensión, las formas de los tensores de entrada deben ser exactamente iguales y se añade una nueva dimensión.
  • torch.cat: Concatena a lo largo de una dimensión existente, los tensores de entrada pueden ser diferentes en la dimensión de concatenación y no se añade una nueva dimensión.

Pytorch torch 参考手册Manual de referencia de Pytorch torch

Otras extensiones