如何使用 PyTorch 仅对矩阵的上三角元素进行 softmax？

Question

给定输入如下：

tensor([[[1.9392, -1.9266,  0.9664],
         [0.0000, -1.9266,  0.9664],
         [0.0000, -0.0000,  0.9664]]])

我想要的输出是：

tensor([[[0.4596,  0.0096, 0.1737],
         [0.0000,  0.0096, 0.1737],
         [0.0000, -0.0000, 0.1737]]])

即仅计算上三角元素的函数。

Answer 1

您可以使用以下命令访问上三角元素torch.triu_indices：

t = tensor([[1.9392, -1.9266, 0.9664], 
            [0.0000, -1.9266, 0.9664], 
            [0.0000, -0.0000, 0.9664]]) 

idx = torch.triu_indices(*t.shape)
soft = F.softmax(t[idx[0], idx[1]], dim=0)

如果您想按照所需的输出重新分配值：

>>> t[idx[0], idx[1]] = soft
>>> t

tensor([[0.4596,  0.0096, 0.1737],
        [0.0000,  0.0096, 0.1737],
        [0.0000, -0.0000, 0.1737]])