Функция torch.stack в PyTorch
Справочное руководство Pytorch torch
torch.stackЭто функция в PyTorch, используемая для объединения нескольких тензоров вдоль нового измерения. Она создает новое измерение и размещает все входные тензоры вдоль этого нового измерения.
Это обычно используется в глубоком обучении в таких сценариях, как создание пакетных данных и объединение выходных данных нескольких моделей.
Определение функции
torch.stack(tensors, dim=0, out=None)
Parameters:
tensors(Sequence of Tensor): Последовательность тензоров для объединения. Все тензоры должны иметь одинаковую форму.dim(int, optional): Измерение, вдоль которого выполняется объединение; по умолчанию 0. Новое измерение будет вставлено в эту позицию.out(Tensor, необязательный): Выходной тензор.
Возвращаемое значение:
torch.Tensor: Возвращает объединённый тензор.
Примеры использования
Пример 1: Базовое объединение
Пример
import torch
# Создайте два одномерных тензора
a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])
# Объединить вдоль новой размерности
c = torch.stack([a, b])
print("Размер a:", a.shape)
print("Размер b:", b.shape)
print("Размер c:", c.shape)
print(c)
# Создайте два одномерных тензора
a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])
# Объединить вдоль новой размерности
c = torch.stack([a, b])
print("Размер a:", a.shape)
print("Размер b:", b.shape)
print("Размер c:", c.shape)
print(c)
Результат вывода:
a 的形状: torch.Size([3])
b 的形状: torch.Size([3])
c 的形状: torch.Size([2, 3])
tensor([[1, 2, 3],
[4, 5, 6]])
Пример 2: Указание размерности объединения
Пример
import torch
a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])
# Объединить вдоль dim=1
c = torch.stack([a, b], dim=1)
print("Размер c:", c.shape)
print(c)
a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])
# Объединить вдоль dim=1
c = torch.stack([a, b], dim=1)
print("Размер c:", c.shape)
print(c)
Результат вывода:
c 的形状: torch.Size([3, 2])
tensor([[1, 4],
[2, 5],
[3, 6]])
Пример 3: Объединение нескольких тензоров
Пример
import torch
# Создать несколько тензоров
tensors = [torch.tensor([i, i+1, i+2]) for i in range(5)]
# Объединить все тензоры
result = torch.stack(tensors)
print("Форма результата:", result.shape)
print(result)
# Создать несколько тензоров
tensors = [torch.tensor([i, i+1, i+2]) for i in range(5)]
# Объединить все тензоры
result = torch.stack(tensors)
print("Форма результата:", result.shape)
print(result)
Результат вывода:
结果形状: torch.Size([5, 3])
tensor([[0, 1, 2],
[1, 2, 3],
[2, 3, 4],
[3, 4, 5],
[4, 5, 6]])
Пример 4: Сохранение нескольких состояний в нейронной сети
Пример
import torch
# Симулировать сохранение значений потерь для нескольких эпох
losses = []
for epoch in range(5):
loss = torch.tensor([epoch * 0.1, (epoch + 1) * 0.1])
losses.append(loss)
# Объединить в стек потери всех эпох
all_losses = torch.stack(losses)
print("Форма потерь по эпохам:", all_losses.shape)
print(all_losses)
# Симулировать сохранение значений потерь для нескольких эпох
losses = []
for epoch in range(5):
loss = torch.tensor([epoch * 0.1, (epoch + 1) * 0.1])
losses.append(loss)
# Объединить в стек потери всех эпох
all_losses = torch.stack(losses)
print("Форма потерь по эпохам:", all_losses.shape)
print(all_losses)
Результат вывода:
各 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]])
Разница между torch.stack и torch.cat
Пример
import torch
a = torch.randn(2, 3)
b = torch.randn(2, 3)
# stack создаёт новое измерение
stack_result = torch.stack([a, b])
print("форма результата stack:", stack_result.shape)
# cat не создаёт новое измерение
cat_result = torch.cat([a, b], dim=0)
print("форма результата cat:", cat_result.shape)
a = torch.randn(2, 3)
b = torch.randn(2, 3)
# stack создаёт новое измерение
stack_result = torch.stack([a, b])
print("форма результата stack:", stack_result.shape)
# cat не создаёт новое измерение
cat_result = torch.cat([a, b], dim=0)
print("форма результата cat:", cat_result.shape)
Результат вывода:
stack 结果形状: torch.Size([2, 2, 3]) cat 结果形状: torch.Size([4, 3])
torch.stack: Стекирует вдоль нового измерения; входные тензоры должны иметь в точности одинаковую форму, и добавляется новое измерение.torch.cat: Конкатенирует вдоль существующего измерения; входные тензоры могут различаться по измерению конкатенации, и новое измерение не добавляется.
Другие расширения