如何使用 PyTorch 仅对矩阵的上三角元素进行软最大化？

问题描述

给定的输入如：

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

我想要的输出是：

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

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

解决方法

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

t = tensor([[1.9392,-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.1737]])