Python PyTorch triu_indices用法及代碼示例

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

用法:

torch.triu_indices(row, col, offset=0, *, dtype=torch.long, device='cpu', layout=torch.strided) → Tensor

參數：

• row(int) -二維矩陣中的行數。

• col(int) -二維矩陣中的列數。

• offset(int) -與主對角線的對角線偏移。默認值：如果未提供，則為 0。

關鍵字參數：

• dtype(torch.dtype， 可選的) -返回張量的所需數據類型。默認值：如果 None ， torch.long 。

• device(torch.device， 可選的) -返回張量的所需設備。默認值：如果 None ，使用當前設備作為默認張量類型(參見 torch.set_default_tensor_type() )。 device 將是 CPU 張量類型的 CPU 和 CUDA 張量類型的當前 CUDA 設備。

• layout(torch.layout， 可選的) -目前僅支持 torch.strided 。

在 2-by-N 張量中返回 row 通過 col 矩陣的上三角部分的索引，其中第一行包含所有索引的行坐標，第二行包含列坐標。索引先按行排序，然後按列排序。

矩陣的上三角部分定義為對角線上和之上的元素。

參數 offset 控製要考慮的對角線。如果offset = 0，則保留主對角線之上和之上的所有元素。正值不包括主對角線上方的對角線，類似地，負值包括主對角線下方的對角線。主對角線是i \in [0, \min\{d_{1}, d_{2}\} - 1] 的索引集\lbrace (i, i) \rbrace，其中d_{1}, d_{2} 是矩陣的維度。

注意

在CUDA上運行時，row * col必須小於2^{59}，以防止計算時溢出。

例子：

>>> a = torch.triu_indices(3, 3)
>>> a
tensor([[0, 0, 0, 1, 1, 2],
        [0, 1, 2, 1, 2, 2]])

>>> a = torch.triu_indices(4, 3, -1)
>>> a
tensor([[0, 0, 0, 1, 1, 1, 2, 2, 3],
        [0, 1, 2, 0, 1, 2, 1, 2, 2]])

>>> a = torch.triu_indices(4, 3, 1)
>>> a
tensor([[0, 0, 1],
        [1, 2, 2]])