PyTorch torch.diagonal 関数

Pytorch torch 参考手册Pytorch torch リファレンスマニュアル

torch.diagonalPyTorch においてテンソルの対角線要素を抽出するための関数です。入力テンソルの指定された対角線のビューを返します。

関数定義

torch.diagonal(input, diagonal=0, dim1=0, dim2=1)

パラメータの説明：

  • input: 入力テンソル
  • diagonal: 対角線インデックス。0 は主対角線を表します。
  • dim1: 最初の次元
  • dim2: 2番目の次元

使用例

例

import torch

# 3x3 行列を作成
x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

# 主対角線を抽出
y = torch.diagonal(x)
print(y)

出力結果は次のとおりです：

tensor([1, 5, 9])

例

import torch

# 3x3 行列を作成
x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

# 主対角線の上側の対角線を抽出
y = torch.diagonal(x, offset=1)
print(y)

出力結果は次のとおりです：

tensor([2, 6])

例

import torch

# 3Dテンソルを作成
x = torch.arange(12).reshape(2, 3, 4)

# 指定された次元の対角線を抽出
y = torch.diagonal(x, offset=0, dim1=1, dim2=2)
print(y.shape)
print(y)

出力結果は次のとおりです：

torch.Size([2, 3])
tensor([[ 0,  5, 10],
        [12, 17, 22]])

Pytorch torch 参考手册Pytorch torch リファレンスマニュアル

その他の拡張