Pytorch: torch.diag()创建对角线张量

torch.diag()

torch.diag是PyTorch中的一个函数，用于从给定的矩阵中提取对角线元素，或者构造一个以给定对角线元素为值的对角矩阵。这个函数对于矩阵分解和转换等操作非常重要。

如果输入是一个向量（1D张量），torch.diag会返回一个以该向量为对角线元素的2D方阵。如果输入是一个矩阵（2D张量），则返回一个包含输入矩阵对角线元素的1D张量。

torch.diag还允许你指定对角线的位置，通过参数diagonal实现。如果diagonal=0，则为主对角线；如果diagonal>0，则为位于主对角线之上的对角线；如果diagonal<0，则为位于主对角线之下的对角线。

语法：

• input (Tensor)：输入张量。
• diagonal (int, optional)：指定的对角线。
• out (Tensor, optional)：输出张量。

举例：

import torch

data = torch.tensor([1,2,3,4])
data_two = torch.diag(data,0)
print(data_two)

结果：

举例二：

import torch

data = torch.tensor(float('inf')).cuda().repeat(3)
data_two = torch.diag(data,0)
print(data_two)

结果: