Python PyTorch diag_embed用法及代碼示例

本文簡要介紹python語言中 torch.diag_embed 的用法。

用法:

torch.diag_embed(input, offset=0, dim1=- 2, dim2=- 1) → Tensor

參數：

• input(Tensor) -輸入張量。必須至少是一維的。

• offset(int,可選的) -考慮哪個對角線。默認值：0(主對角線)。

• dim1(int,可選的) -對角線的第一個維度。默認值：-2。

• dim2(int,可選的) -相對於其取對角線的第二維。默認值：-1。

創建一個張量，其某些 2D 平麵(由 dim1 和 dim2 指定)的對角線由 input 填充。為了便於創建批量對角矩陣，默認選擇由返回張量的最後兩個維度形成的 2D 平麵。

參數 offset 控製要考慮的對角線：

• 如果offset = 0，它是主對角線。

• 如果offset > 0，它在主對角線之上。

• 如果offset < 0，則它位於主對角線下方。

新矩陣的大小將被計算為最後一個輸入維度的大小的指定對角線。請注意，對於 offset 以外的 0 ， dim1 和 dim2 的順序很重要。交換它們相當於更改 offset 的符號。

使用相同的參數將 torch.diagonal() 應用於此函數的輸出會產生與輸入相同的矩陣。但是， torch.diagonal() 具有不同的默認尺寸，因此需要明確指定。

例子：

>>> a = torch.randn(2, 3)
>>> torch.diag_embed(a)
tensor([[[ 1.5410,  0.0000,  0.0000],
         [ 0.0000, -0.2934,  0.0000],
         [ 0.0000,  0.0000, -2.1788]],

        [[ 0.5684,  0.0000,  0.0000],
         [ 0.0000, -1.0845,  0.0000],
         [ 0.0000,  0.0000, -1.3986]]])

>>> torch.diag_embed(a, offset=1, dim1=0, dim2=2)
tensor([[[ 0.0000,  1.5410,  0.0000,  0.0000],
         [ 0.0000,  0.5684,  0.0000,  0.0000]],

        [[ 0.0000,  0.0000, -0.2934,  0.0000],
         [ 0.0000,  0.0000, -1.0845,  0.0000]],

        [[ 0.0000,  0.0000,  0.0000, -2.1788],
         [ 0.0000,  0.0000,  0.0000, -1.3986]],

        [[ 0.0000,  0.0000,  0.0000,  0.0000],
         [ 0.0000,  0.0000,  0.0000,  0.0000]]])