Функция torch.diagonal в PyTorch

Pytorch torch 参考手册Справочное руководство Pytorch torch

torch.diagonalЭто функция в PyTorch для извлечения диагональных элементов тензора. Она возвращает представление указанной диагонали входного тензора.

Определение функции

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

Описание параметров:

  • input: Входной тензор
  • diagonal: Индекс диагонали, 0 соответствует главной диагонали
  • dim1: Первое измерение
  • dim2: Второе измерение

Примеры использования

Пример

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

Другие расширения