PyTorchでテンソルの要素を並べ替える方法【torch.sort()の使い方】

PyTorchでテンソルの要素を並べ替えるには、torch.sort() メソッドを使用します。このメソッドは2つのテンソルを返します。1つ目は要素がソートされた値を持つテンソル、2つ目は元のテンソルにおける各要素のインデックス（位置情報）を持つテンソルです。また、2次元テンソルに対しては、行方向・列方向それぞれに沿ってソートを計算することもできます。

手順

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

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

• 作成したテンソルの要素をソートするために、torch.sort(input, dim) を計算し、その結果を新しい変数 v に代入します。ここで input は入力テンソル、dim はソートを行う軸（次元）を表します。行方向にソートする場合は dim を 1 に、列方向にソートする場合は dim を 0 に設定します。

• ソート済みの値を持つテンソルは v[0] で、ソートされた要素の元のインデックスを持つテンソルは v[1] でアクセスできます。

• ソートされた値のテンソルと、そのインデックスのテンソルを出力します。

例1：1次元テンソルのソート

次のPythonプログラムは、1次元テンソルの要素をソートする方法を示しています。

# テンソルの要素をソートするPythonプログラム
# 必要なライブラリをインポート
import torch

# テンソルを作成
T = torch.Tensor([2.334,4.433,-4.33,-0.433,5, 4.443])
print("Original Tensor:\n", T)

# テンソルTをソート
# 昇順にソートされる
v = torch.sort(T)

# ソート済みの値を持つテンソルを出力
print("Tensor with sorted value:\n", v[0])

# ソート済みの値のインデックスを出力
print("Indices of sorted value:\n", v[1])

出力結果

Original Tensor:
    tensor([ 2.3340, 4.4330, -4.3300, -0.4330, 5.0000, 4.4430])
Tensor with sorted value:
    tensor([-4.3300, -0.4330, 2.3340, 4.4330, 4.4430, 5.0000])
Indices of sorted value:
    tensor([2, 3, 0, 1, 5, 4])

例2：2次元テンソルのソート

次のPythonプログラムは、2次元テンソルの要素をソートする方法を示しています。デフォルトでは最後の次元（行方向）に沿ってソートされますが、dim引数を指定することで列方向や行方向のソートを明示的に切り替えられます。

# 2次元テンソルの要素をソートするPythonプログラム
# ライブラリをインポート
import torch

# 2次元テンソルを作成
T = torch.Tensor([[2,3,-32],
                  [43,4,-53],
                  [4,37,-4],
                  [3,-75,34]])
print("Original Tensor:\n", T)

# テンソルTをソート
# 昇順にソートされる
v = torch.sort(T)

# ソート済みの値を持つテンソルを出力
print("Tensor with sorted value:\n", v[0])

# ソート済みの値のインデックスを出力
print("Indices of sorted value:\n", v[1])
print("Sort tensor Column-wise")
v = torch.sort(T, 0)

# 列方向にソートした値のテンソルを出力
print("Tensor with sorted value:\n", v[0])

# インデックスを出力
print("Indices of sorted value:\n", v[1])
print("Sort tensor Row-wise")
v = torch.sort(T, 1)

# 行方向にソートした値のテンソルを出力
print("Tensor with sorted value:\n", v[0])

# インデックスを出力
print("Indices of sorted value:\n", v[1])

出力結果

Original Tensor:
tensor([[ 2., 3., -32.],
        [ 43., 4., -53.],
        [ 4., 37., -4.],
        [ 3., -75., 34.]])
Tensor with sorted value:
tensor([[-32., 2., 3.],
        [-53., 4., 43.],
        [ -4., 4., 37.],
        [-75., 3., 34.]])
Indices of sorted value:
tensor([[2, 0, 1],
        [2, 1, 0],
        [2, 0, 1],
        [1, 0, 2]])
Sort tensor Column-wise
Tensor with sorted value:
tensor([[ 2., -75., -53.],
        [ 3., 3., -32.],
        [ 4., 4., -4.],
        [ 43., 37., 34.]])
Indices of sorted value:
tensor([[0, 3, 1],
        [3, 0, 0],
        [2, 1, 2],
        [1, 2, 3]])
Sort tensor Row-wise
Tensor with sorted value:
tensor([[-32., 2., 3.],
        [-53., 4., 43.],
        [ -4., 4., 37.],
        [-75., 3., 34.]])
Indices of sorted value:
tensor([[2, 0, 1],
        [2, 1, 0],
        [2, 0, 1],
        [1, 0, 2]])

まとめ

PyTorchでテンソルをソートする際は、torch.sort() を使うことで「ソート済みの値」と「元の位置を示すインデックス」の両方を一度に取得できます。dim引数を指定すれば、2次元以上のテンソルでも行方向（dim=1）や列方向（dim=0）など、任意の軸に沿った柔軟なソートが可能です。