PyTorchでテンソルをスクイーズ・アンスクイーズする方法を解説

PyTorchでは、テンソルのサイズが1の次元を削除することを「スクイーズ（squeeze）」、逆にサイズ1の新しい次元を挿入することを「アンスクイーズ（unsqueeze）」と呼びます。この記事では、それぞれの操作を行うためのtorch.squeeze()メソッドとtorch.unsqueeze()メソッドの使い方を、具体的なコード例とともに解説します。

torch.squeeze() とは

torch.squeeze()は、入力テンソルのすべての次元のうち、サイズが1の次元を取り除いた新しいテンソルを返します。例えば、入力テンソルの形状が (M × 1 × N × 1 × P) の場合、スクイーズ後のテンソルの形状は (M × N × P) になります。

torch.unsqueeze() とは

torch.unsqueeze()は、指定した位置にサイズ1の新しい次元を挿入したテンソルを返します。バッチ次元の追加など、テンソルの形状を調整したい場面でよく使われます。

実行手順

• 必要なライブラリをインポートします。以下のPythonサンプルでは torch ライブラリを使用します。事前にインストールしておいてください。

• テンソルを作成し、内容を出力します。

• torch.squeeze(input) を実行します。サイズ1の次元が削除され、input テンソルのその他の次元を持つテンソルが返されます。

• torch.unsqueeze(input, dim) を実行します。指定した dim の位置にサイズ1の新しい次元が挿入されたテンソルが返されます。

• スクイーズ／アンスクイーズ後のテンソルを出力して確認します。

例1：テンソルのスクイーズ

# Python program to squeeze and unsqueeze a tensor
# import necessary library
import torch

# Create a tensor of all one
T = torch.ones(2,1,2) # size 2x1x2
print("Original Tensor T:\n", T )
print("Size of T:", T.size())

# Squeeze the dimension of the tensor
squeezed_T = torch.squeeze(T) # now size 2x2
print("Squeezed_T\n:", squeezed_T )
print("Size of Squeezed_T:", squeezed_T.size())

出力結果

Original Tensor T:
tensor([[[1., 1.]],
        [[1., 1.]]])
Size of T: torch.Size([2, 1, 2])
Squeezed_T
: tensor([[1., 1.],
        [1., 1.]])
Size of Squeezed_T: torch.Size([2, 2])

この例では、形状 (2, 1, 2) のテンソルに対して torch.squeeze() を適用すると、サイズ1の次元が削除され、形状が (2, 2) に変化していることがわかります。

例2：テンソルのアンスクイーズ

# Python program to squeeze and unsqueeze a tensor
# import necessary library
import torch

# create a tensor
T = torch.Tensor([1,2,3]) # size 3
print("Original Tensor T:\n", T )
print("Size of T:", T.size())

# Squeeze the tensor in dimension o or column dim
unsqueezed_T = torch.unsqueeze(T, dim = 0) # now size 1x3
print("Unsqueezed T\n:", unsqueezed_T )
print("Size of UnSqueezed T:", unsqueezed_T.size())

# Squeeze the tensor in dimension 1 or row dim
unsqueezed_T = torch.unsqueeze(T, dim = 1) # now size 3x1
print("Unsqueezed T\n:", unsqueezed_T )
print("Size of Unsqueezed T:", unsqueezed_T.size())

出力結果

Original Tensor T:
    tensor([1., 2., 3.])
Size of T: torch.Size([3])
Unsqueezed T
: tensor([[1., 2., 3.]])
Size of UnSqueezed T: torch.Size([1, 3])
Unsqueezed T
: tensor([[1.],
        [2.],
        [3.]])
Size of Unsqueezed T: torch.Size([3, 1])

この例では、サイズ3の1次元テンソルに対して dim=0 を指定すると形状が (1, 3) になり、dim=1 を指定すると形状が (3, 1) になります。このように、挿入する位置によってテンソルの形状がどのように変わるかを確認できます。

まとめ

torch.squeeze()はサイズ1の次元を削除し、torch.unsqueeze()は指定した位置にサイズ1の次元を挿入します。両者はニューラルネットワークの入力データの整形やバッチ処理などで頻繁に使われる重要な操作なので、しっかり理解しておきましょう。