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])