Python PyTorch diag用法及代码示例

本文简要介绍python语言中 torch.diag 的用法。

用法:

torch.diag(input, diagonal=0, *, out=None) → Tensor

参数：

• input(Tensor) -输入张量。

• diagonal(int,可选的) -要考虑的对角线

关键字参数：

out(Tensor,可选的) -输出张量。

• 如果input 是一个向量(一维张量)，则返回一个以input 的元素为对角线的二维平方张量。

• 如果 input 是矩阵(二维张量)，则返回具有 input 的对角元素的一维张量。

参数 diagonal 控制要考虑的对角线：

• 如果 diagonal = 0，则为主对角线。

• 如果 diagonal > 0，则位于主对角线上方。

• 如果 diagonal < 0，则它位于主对角线下方。

例子：

获取输入向量为对角线的方阵：

>>> a = torch.randn(3)
>>> a
tensor([ 0.5950,-0.0872, 2.3298])
>>> torch.diag(a)
tensor([[ 0.5950, 0.0000, 0.0000],
        [ 0.0000,-0.0872, 0.0000],
        [ 0.0000, 0.0000, 2.3298]])
>>> torch.diag(a, 1)
tensor([[ 0.0000, 0.5950, 0.0000, 0.0000],
        [ 0.0000, 0.0000,-0.0872, 0.0000],
        [ 0.0000, 0.0000, 0.0000, 2.3298],
        [ 0.0000, 0.0000, 0.0000, 0.0000]])

获取给定矩阵的 k-th 对角线：

>>> a = torch.randn(3, 3)
>>> a
tensor([[-0.4264, 0.0255,-0.1064],
        [ 0.8795,-0.2429, 0.1374],
        [ 0.1029,-0.6482,-1.6300]])
>>> torch.diag(a, 0)
tensor([-0.4264,-0.2429,-1.6300])
>>> torch.diag(a, 1)
tensor([ 0.0255, 0.1374])