Fonction torch.diagonal de PyTorch

Pytorch torch 参考手册Manuel de référence Pytorch torch

torch.diagonalIl s'agit d'une fonction de PyTorch pour extraire les éléments diagonaux d'un tenseur. Elle renvoie une vue de la diagonale spécifiée du tenseur d'entrée.

Définition de fonction

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

Description des paramètres :

  • input: tenseur d’entrée
  • diagonal: Indice de la diagonale, 0 représente la diagonale principale
  • dim1: Première dimension
  • dim2: Deuxième dimension

Exemples d'utilisation

Exemple

import torch

# Créer une matrice 3x3
x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

# Extraire la diagonale principale
y = torch.diagonal(x)
print(y)

Le résultat de sortie est:

tensor([1, 5, 9])

Exemple

import torch

# Créer une matrice 3x3
x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

# Extraire la diagonale au-dessus de la diagonale principale
y = torch.diagonal(x, offset=1)
print(y)

Le résultat de sortie est:

tensor([2, 6])

Exemple

import torch

# Créer un tenseur 3D
x = torch.arange(12).reshape(2, 3, 4)

# Extraire la diagonale des dimensions spécifiées
y = torch.diagonal(x, offset=0, dim1=1, dim2=2)
print(y.shape)
print(y)

Le résultat de sortie est:

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

Pytorch torch 参考手册Manuel de référence Pytorch torch

Autres extensions