介绍
torch.triu() 函数，用于提取矩阵的 上三角部分，是张量操作中很常用的一个工具。
作用
返回一个与输入张量形状相同的新张量，其中 上三角部分（包括指定对角线）保留原值，其余位置填充为 0。
上三角矩阵
矩阵中：
- 行号 > 列号 的位置为0（即：矩阵的 左下是0 ）
- 行号 ≤ 列号 的位置为有效元素（即：矩阵的 右上是有效元素 ）
原始矩阵：
[[1., 1., 1., 1.],
[1., 1., 1., 1.],
[1., 1., 1., 1.],
[1., 1., 1., 1.]]
上三角矩阵（diagonal=0）：
[[1., 1., 1., 1.],
[0., 1., 1., 1.],
[0., 0., 1., 1.],
[0., 0., 0., 1.]]
语法
torch.triu(input, diagonal=0, *, out=None) → Tensor
参数：
input：输入的张量（通常是 2D 矩阵，也支持更高维张量，会对最后两个维度操作）。diagonal：指定对角线位置，默认值为 0：diagonal=0：主对角线及以上保留（标准上三角）。diagonal>0：主对角线上方的对角线（如diagonal=1表示主对角线右侧第一条对角线及以上）。diagonal<0：主对角线下方的对角线（如diagonal=-1表示主对角线左侧第一条对角线及以上）。
例子
下面通过具体例子帮你理解不同参数的效果：
import torch
# 1. 基础用法（2D矩阵，默认diagonal=0）
mat = torch.tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
triu_0 = torch.triu(mat)
print("diagonal=0（主对角线及以上）：")
print(triu_0)
# 输出：
# tensor([[1, 2, 3],
# [0, 5, 6],
# [0, 0, 9]])
# 2. diagonal=1（主对角线上方第一条对角线及以上）
triu_1 = torch.triu(mat, diagonal=1)
print("\ndiagonal=1：")
print(triu_1)
# 输出：
# tensor([[0, 2, 3],
# [0, 0, 6],
# [0, 0, 0]])
# 3. diagonal=-1（主对角线左侧第一条对角线及以上）
triu_neg1 = torch.triu(mat, diagonal=-1)
print("\ndiagonal=-1：")
print(triu_neg1)
# 输出：
# tensor([[1, 2, 3],
# [4, 5, 6],
# [0, 8, 9]])
# 4. 高维张量（对最后两个维度操作）
high_dim_mat = torch.tensor([[[1,2],[3,4]], [[5,6],[7,8]]])
triu_high_dim = torch.triu(high_dim_mat)
print("\n高维张量操作结果：")
print(triu_high_dim)
# 输出：
# tensor([[[1, 2],
# [0, 4]],
# [[5, 6],
# [0, 8]]])
常见使用场景
- 注意力机制：在 Transformer 的自注意力计算中，常用
torch.triu()生成掩码（mask），遮挡未来位置的信息（如torch.triu(torch.ones(size), diagonal=1)生成上三角全 1 的掩码，再取反得到下三角掩码）。 - 矩阵运算：仅需要矩阵上三角部分参与计算时（如对称矩阵的上三角提取）。
总结
torch.triu()核心功能是提取张量的上三角部分，非上三角区域置 0；diagonal参数控制保留的对角线位置，0 是主对角线，正数向上偏移，负数向下偏移；- 支持高维张量，默认对最后两个维度（矩阵维度）操作，是注意力掩码等场景的核心工具。