使用 Softmax 后如何减少张量的维数？

Question

我得到了一个分数张量（我们称之为logits_tensor ），其形状为： (1910, 164, 33) 。

看看它， logits_tensor[0][0] ：

tensor([-2.5916, -1.5290, -0.8218, -0.8882, -2.0961, -2.1064, -0.7842, -1.5200,
        -2.1324, -1.5561, -2.4731, -2.1933, -2.8489, -1.8257, -1.8033, -1.8771,
        -2.8365,  0.6690, -0.6895, -1.7054, -2.4862, -0.8104, -1.5395, -1.1351,
        -2.7154, -1.7646, -2.6595, -2.0591, -2.7554, -1.8661, -2.7512, -2.0655,
         5.7374])

现在，通过应用softmax

probs_tensor = torch.nn.functional.softmax(logits_tensor, dim=-1)

我获得了另一个包含概率的具有相同维度的张量probs_tensor[0][0] ：

tensor([2.3554e-04, 6.8166e-04, 1.3825e-03, 1.2937e-03, 3.8660e-04, 3.8263e-04,
        1.4356e-03, 6.8778e-04, 3.7283e-04, 6.6341e-04, 2.6517e-04, 3.5078e-04,
        1.8211e-04, 5.0665e-04, 5.1810e-04, 4.8127e-04, 1.8438e-04, 6.1396e-03,
        1.5782e-03, 5.7138e-04, 2.6173e-04, 1.3984e-03, 6.7454e-04, 1.0107e-03,
        2.0812e-04, 5.3857e-04, 2.2009e-04, 4.0118e-04, 1.9996e-04, 4.8660e-04,
        2.0079e-04, 3.9860e-04, 9.7570e-01])

我想要获得的是一个形状为1910, 164)的张量，其中包含上面显示的最大概率（对于 164 个元素中的每一个）的索引，如下所示：

precitions[0]
> tensor([32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32,  1, 17, 17, 17,
       17, 17, 17, 17, 17, 17, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32,
       32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32,  0,  0,  0,
        0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
        0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
        0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
        0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
        0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
        0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
        0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0]

请注意，“32”是probs_tensor[0][0]中较高概率元素的索引。 使用torch.argmax可以完成相同的任务，但我需要 softmax 步骤。

Answer 1

实际上，您可以在张量上应用torch.argmax ：

>>> logits_tensor = torch.rand(1910, 164, 33)
>>> probs_tensor = logits_tensor.softmax(-1)

>>> probs_tensor.argmax(-1).shape
torch.Size([1910, 164])

请注意，在logits_tensor上应用 argmax 与在probs_tensor上应用它是相同的。 具有最高值的 logit 将保持具有最高概率质量的 logit。