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)
结果: