PyTorch torch.stack 関数
Pytorch torch リファレンスマニュアルtorch.stackPyTorchで新しい次元に沿って複数のテンソルをスタックするための関数です。新しい次元を作成し、すべての入力テンソルをこの新しい次元に配置します。
これは深層学習で、バッチデータの作成や複数モデルの出力のスタックなどの場面でよく使用されます。
関数定義
torch.stack(tensors, dim=0, out=None)
パラメータ:
tensors(Sequence of Tensor): スタックするテンソルのシーケンス。すべてのテンソルは同じ形状である必要があります。dim(int, オプション): スタックする次元。デフォルトは0です。新しい次元はこの位置に挿入されます。out(Tensor, オプション): 出力テンソル。
戻り値:
torch.Tensor: スタックされたテンソルを返します。
使用例
例1: 基本スタック
実例
import torch
# 2つの一次元テンソルを作成
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)
# 2つの一次元テンソルを作成
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: 新しい次元に沿ってスタックします。入力テンソルの形状は完全に同じである必要があり、新しい次元が1つ増えます。torch.cat: 既存の次元に沿って連結します。入力テンソルは連結次元で異なっていてもよく、新しい次元は増えません。
その他の拡張