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