PyTorchでテンソルを結合する方法を徹底解説！torch.cat()とtorch.stack()の違いと使い分け

PyTorchでは、torch.cat()とtorch.stack()の2つの関数を使って、2つ以上のテンソルを結合できます。torch.cat()は既存の次元に沿ってテンソルを連結（concatenate）する関数、torch.stack()は新しい次元を作成してテンソルを積み重ねる（stack）関数です。どちらも0次元や-1次元など、任意の次元を指定してテンソルを結合することが可能です。

一見すると似た機能を持つこの2つのメソッドですが、基本的な違いは何なのでしょうか？

• torch.cat()：既存の次元に沿ってテンソルのシーケンスを連結します。そのため、テンソルの次元数は変化しません。

• torch.stack()：新しい次元に沿ってテンソルを積み重ねます。その結果、出力されるテンソルの次元数が1つ増加します。

実装の手順

• 必要なライブラリをインポートします。以下のすべての例で使用するPythonライブラリはtorchです。事前にインストール済みであることを確認してください。

• 2つ以上のPyTorchテンソルを作成し、内容を出力して確認します。

• torch.cat()またはtorch.stack()を使用して、作成したテンソルを結合します。引数として次元（0や-1など）を指定することで、特定の次元に沿って結合できます。

• 最後に、連結またはスタックされた結果のテンソルを出力します。

例1：torch.cat()で1次元テンソルを連結する

# PyTorchでテンソルを結合するプログラム
# 必要なライブラリをインポート
import torch

# テンソルを作成
T1 = torch.Tensor([1,2,3,4])
T2 = torch.Tensor([0,3,4,1])
T3 = torch.Tensor([4,3,2,5])

# 作成したテンソルを出力
print("T1:", T1)
print("T2:", T2)
print("T3:", T3)

# torch.cat()でテンソルを連結
T = torch.cat((T1,T2,T3))
# 連結後のテンソルを出力
print("T:",T)

出力結果

上記のPythonコードを実行すると、以下の出力が得られます。

T1: tensor([1., 2., 3., 4.])
T2: tensor([0., 3., 4., 1.])
T3: tensor([4., 3., 2., 5.])
T: tensor([1., 2., 3., 4., 0., 3., 4., 1., 4., 3., 2., 5.])

デフォルトでは0次元（行方向）に沿って連結されるため、3つの1次元テンソルが1本の長い1次元テンソルにつながっていることがわかります。

例2：torch.cat()で2次元テンソルを連結する

# 必要なライブラリをインポート
import torch

# テンソルを作成
T1 = torch.Tensor([[1,2],[3,4]])
T2 = torch.Tensor([[0,3],[4,1]])
T3 = torch.Tensor([[4,3],[2,5]])

# 作成したテンソルを出力
print("T1:\n", T1)
print("T2:\n", T2)
print("T3:\n", T3)

print("0次元でテンソルを連結")
T = torch.cat((T1,T2,T3), 0)
print("T:\n", T)

print("-1次元でテンソルを連結")
T = torch.cat((T1,T2,T3), -1)
print("T:\n", T)

出力結果

上記のPythonコードを実行すると、以下の出力が得られます。

T1:
tensor([[1., 2.],
        [3., 4.]])
T2:
tensor([[0., 3.],
        [4., 1.]])
T3:
tensor([[4., 3.],
        [2., 5.]])
0次元でテンソルを連結
T:
tensor([[1., 2.],
        [3., 4.],
        [0., 3.],
        [4., 1.],
        [4., 3.],
        [2., 5.]])
-1次元でテンソルを連結
T:
tensor([[1., 2., 0., 3., 4., 3.],
        [3., 4., 4., 1., 2., 5.]])

上記の例では、2次元テンソルを0次元と-1次元に沿って連結しています。0次元で連結すると行数が増え、列数は変化しません。一方、-1次元（最後の次元＝列方向）で連結すると列数が増え、行数は変化しないことが確認できます。

例3：torch.stack()で1次元テンソルを積み重ねる

# PyTorchでテンソルを結合するプログラム
# 必要なライブラリをインポート
import torch

# テンソルを作成
T1 = torch.Tensor([1,2,3,4])
T2 = torch.Tensor([0,3,4,1])
T3 = torch.Tensor([4,3,2,5])

# 作成したテンソルを出力
print("T1:", T1)
print("T2:", T2)
print("T3:", T3)

# torch.stack()でテンソルを積み重ねる
print("テンソルをスタック")
T = torch.stack((T1,T2,T3))

# 結合後のテンソルを出力
print("T:\n",T)
print("0次元でテンソルをスタック")
T = torch.stack((T1,T2,T3), 0)

print("T:\n", T)
print("-1次元でテンソルをスタック")
T = torch.stack((T1,T2,T3), -1)
print("T:\n", T)

出力結果

上記のPythonコードを実行すると、以下の出力が得られます。

T1: tensor([1., 2., 3., 4.])
T2: tensor([0., 3., 4., 1.])
T3: tensor([4., 3., 2., 5.])
テンソルをスタック
T:
tensor([[1., 2., 3., 4.],
        [0., 3., 4., 1.],
        [4., 3., 2., 5.]])
0次元でテンソルをスタック
T:
tensor([[1., 2., 3., 4.],
        [0., 3., 4., 1.],
        [4., 3., 2., 5.]])
-1次元でテンソルをスタック
T:
tensor([[1., 0., 4.],
        [2., 3., 3.],
        [3., 4., 2.],
        [4., 1., 5.]])

上記の例からわかるように、1次元テンソルをスタックすると、結果は2次元テンソルになります。これがtorch.cat()との大きな違いです。

例4：torch.stack()で2次元テンソルを積み重ねる

# 必要なライブラリをインポート
import torch

# テンソルを作成
T1 = torch.Tensor([[1,2],[3,4]])
T2 = torch.Tensor([[0,3],[4,1]])
T3 = torch.Tensor([[4,3],[2,5]])

# 作成したテンソルを出力
print("T1:\n", T1)
print("T2:\n", T2)
print("T3:\n", T3)

print("0次元でテンソルをスタック")
T = torch.stack((T1,T2,T3), 0)
print("T:\n", T)

print("-1次元でテンソルをスタック")
T = torch.stack((T1,T2,T3), -1)
print("T:\n", T)

出力結果

上記のPythonコードを実行すると、以下の出力が得られます。

T1:
tensor([[1., 2.],
        [3., 4.]])
T2:
tensor([[0., 3.],
        [4., 1.]])
T3:
tensor([[4., 3.],
        [2., 5.]])
0次元でテンソルをスタック
T:
tensor([[[1., 2.],
         [3., 4.]],
         [[0., 3.],
         [4., 1.]],
         [[4., 3.],
         [2., 5.]]])
-1次元でテンソルをスタック
T:
tensor([[[1., 0., 4.],
         [2., 3., 3.]],
         [[3., 4., 2.],
         [4., 1., 5.]]])

上記の例では、2次元テンソル同士をスタックすることで、3次元テンソルが生成されていることが確認できます。

まとめ

• torch.cat()は既存の次元に沿ってテンソルを連結し、次元数は変化しません。

• torch.stack()は新しい次元を作ってテンソルを積み重ねるため、次元数が1つ増えます。

• どちらも第2引数でdim（次元）を指定でき、0は先頭の次元、-1は最後の次元を意味します。

バッチ処理で複数のテンソルをまとめたい場合はtorch.stack()、単純にデータをつなげたい場合はtorch.cat()というように、目的に応じて使い分けるとよいでしょう。