10分で理解！diag() 関数による PyTorch Tensor の対角線操作

PyTorch Tensor の torch.Tensor.diag() 関数

入力 Tensor が 1次元の場合、torch.Tensor.diag() はその要素を対角線に並べた 2次元正方行列を返します。

例：

import torch

# 1次元 Tensor
input_tensor = torch.tensor([1, 2, 3])

# 対角線行列に変換
output_tensor = torch.diag(input_tensor)

print(output_tensor)
# tensor([[1, 0, 0],
#        [0, 2, 0],
#        [0, 0, 3]])

入力 Tensor が 2次元正方行列の場合、torch.Tensor.diag() はその対角線の要素を 1次元 Tensor として返します。

例：

# 2次元正方行列
input_tensor = torch.tensor([[1, 2, 3],
                             [4, 5, 6],
                             [7, 8, 9]])

# 対角線の要素を取得
output_tensor = torch.diag(input_tensor)

print(output_tensor)
# tensor([1, 5, 9])

オプション引数

torch.Tensor.diag() 関数には、以下のオプション引数があります。

• diagonal: 対角線に配置する要素を指定する Tensor です。デフォルトは入力 Tensor 自身です。
• offset: 対角線の位置をオフセットする整数です。デフォルトは 0 です。
• dim1: 対角線要素を取り出す次元を指定する整数です。デフォルトは 0 です。

これらのオプション引数を使うことで、より複雑な操作を行うことができます。

関連関数

torch.Tensor.diag() 関数と関連する関数として、以下のものがあります。

• torch.diag_embed(): 対角線要素を指定して 2次元正方行列を作成する関数です。
• torch.tril(): 下三角行列を取得する関数です。

torch.Tensor.diag() 関数は、PyTorch の Tensor クラスに属する関数で、入力 Tensor の次元に応じて異なる動作を持ちます。オプション引数や関連関数を使うことで、様々な操作を行うことができます。

PyTorch Tensor.diag() のサンプルコード

1次元 Tensor を対角線行列に変換する

import torch

# 1次元 Tensor
input_tensor = torch.tensor([1, 2, 3, 4, 5])

# 対角線行列に変換
output_tensor = torch.diag(input_tensor)

print(output_tensor)
# tensor([[1, 0, 0, 0, 0],
#        [0, 2, 0, 0, 0],
#        [0, 0, 3, 0, 0],
#        [0, 0, 0, 4, 0],
#        [0, 0, 0, 0, 5]])

2次元正方行列の対角線の要素を取得する

# 2次元正方行列
input_tensor = torch.tensor([[1, 2, 3],
                             [4, 5, 6],
                             [7, 8, 9]])

# 対角線の要素を取得
output_tensor = torch.diag(input_tensor)

print(output_tensor)
# tensor([1, 5, 9])

対角線要素を指定して 2次元正方行列を作成する

import torch

# 対角線要素
diagonal_elements = torch.tensor([1, 2, 3])

# 2次元正方行列を作成
output_tensor = torch.diag_embed(diagonal_elements)

print(output_tensor)
# tensor([[1, 0, 0],
#        [0, 2, 0],
#        [0, 0, 3]])

対角線要素をオフセットする

# 1次元 Tensor
input_tensor = torch.tensor([1, 2, 3])

# 対角線行列に変換 (オフセット = 1)
output_tensor = torch.diag(input_tensor, offset=1)

print(output_tensor)
# tensor([[0, 1, 0, 0, 0],
#        [0, 0, 2, 0, 0],
#        [0, 0, 0, 3, 0],
#        [0, 0, 0, 0, 0],
#        [0, 0, 0, 0, 0]])

複数次元 Tensor の特定次元から対角線要素を取得する

# 3次元 Tensor
input_tensor = torch.tensor([[[1, 2, 3],
                              [4, 5, 6]],
                             [[7, 8, 9],
                              [10, 11, 12]]])

# 次元 1 と 2 から対角線要素を取得
output_tensor = torch.diag(input_tensor, dim1=1, dim2=2)

print(output_tensor)
# tensor([1, 5, 10])

下三角行列を取得する

# 2次元正方行列
input_tensor = torch.tensor([[1, 2, 3],
                             [4, 5, 6],
                             [7, 8, 9]])

# 下三角行列を取得
output_tensor = torch.tril(input_tensor)

print(output_tensor)
# tensor([[1, 0, 0],
#        [4, 5, 0],
#        [7, 8, 9]])

上三角行列を取得する

# 2次元正方行列
input_tensor = torch.tensor([[1, 2, 3],
                             [4, 5, 6],
                             [7, 8, 9]])

# 上三角行列を取得
output_tensor = torch.triu(input_tensor)

print(output_tensor)
# tensor([[1, 2, 3],
#        [0, 5, 6],
#        [0, 0, 9]])

Tensor.diag() 関数の代替方法

ループ処理

forループを使って、手動で対角線要素を抽出したり、対角線行列を作成することができます。

import torch

# 1次元 Tensor
input_tensor = torch.tensor([1, 2, 3])

# 対角線要素を抽出
diagonal_elements = []
for i in range(input_tensor.size(0)):
    diagonal_elements.append(input_tensor[i, i])

# 対角線行列を作成
output_tensor = torch.zeros(input_tensor.size(0), input_tensor.size(0))
for i in range(input_tensor.size(0)):
    output_tensor[i, i] = diagonal_elements[i]

print(output_tensor)
# tensor([[1, 0, 0],
#        [0, 2, 0],
#        [0, 0, 3]])

スライス操作を使って、対角線要素を抽出したり、対角線行列を作成することができます。

# 2次元正方行列
input_tensor = torch.tensor([[1, 2, 3],
                             [4, 5, 6],
                             [7, 8, 9]])

# 対角線要素を抽出
diagonal_elements = input_tensor.diagonal()

# 対角線行列を作成
output_tensor = torch.diag(diagonal_elements)

print(output_tensor)
# tensor([[1, 0, 0],
#        [0, 5, 0],
#        [0, 0, 9]])

view() 関数を使って、Tensor の形状を変換することで、対角線要素を抽出したり、対角線行列を作成することができます。

# 1次元 Tensor
input_tensor = torch.tensor([1, 2, 3, 4, 5])

# 対角線要素を抽出
diagonal_elements = input_tensor.view(-1, 1)

# 対角線行列を作成
output_tensor = input_tensor.view(5, 5).diag()

print(output_tensor)
# tensor([[1, 0, 0, 0, 0],
#        [0, 2, 0, 0, 0],
#        [0, 0, 3, 0, 0],
#        [0, 0, 0, 4, 0],
#        [0, 0, 0, 0, 5]])

その他のライブラリ

NumPy や SciPy などのライブラリには、対角線行列の操作に特化した関数があります。

import numpy as np

# 1次元 Tensor を NumPy 配列に変換
input_array = input_tensor.numpy()

# 対角線行列を作成
output_array = np.diag(input_array)

# NumPy 配列を PyTorch Tensor に変換
output_tensor = torch.from_numpy(output_array)

print(output_tensor)
# tensor([[1, 0, 0, 0, 0],
#        [0, 2, 0, 0, 0],
#        [0, 0, 3, 0, 0],
#        [0, 0, 0, 4, 0],
#        [0, 0, 0, 0, 5]])

Tensor.diag() 関数は、対角線行列の操作に便利な関数ですが、上記のような代替方法もあります。状況に応じて、最適な方法を選択してください。