使用 Softmax 后如何减少张量的维度？ [火炬]

Question

我得到了一个分数张量（我们称之为 logits_tensor），其形状为：torch.Size([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])

我想要获得的是一个形状为 torch.Size([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-step。你可以帮帮我吗？谢谢！

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

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

使用 Softmax 后如何减少张量的维度？ [火炬]

How can I reduce dimension of a tensor after using Softmax? [PyTorch]

python

pytorch

tensor

torch

softmax