PyTorch神经网络：将预测概率转为长度为4的预测结果张量

你的输出是4×6的logits张量（每个样本对应6个类别的原始模型输出，未转换为概率），要得到和标签格式一致的长度为4的类别张量，可按以下方式处理：

1. 直接获取预测类别（无需概率） #

Softmax是单调递增函数，logits的最大值位置与Softmax后的最大值位置完全一致，因此可直接对logits使用argmax函数，指定维度为1（按行取最大值索引）：

import torch

# 假设你的输出张量为output
output = torch.tensor([[-0.2713, -0.6608, -0.4430, -0.0207, -0.4408, -0.3075],
                       [-0.2713, -0.6608, -0.4430, -0.0207, -0.4408, -0.3075],
                       [-0.2713, -0.6608, -0.4430, -0.0207, -0.4408, -0.3075],
                       [-0.2713, -0.6608, -0.4430, -0.0207, -0.4408, -0.3075]])

# 获取每个样本的预测类别索引
preds = output.argmax(dim=1)
print(preds)
# 输出：tensor([3, 3, 3, 3])

得到的preds就是长度为4的张量，与标签格式完全匹配。

2. 先转换为概率再获取类别 #

如果需要先得到每个类别的概率分布，再取对应类别，可先用Softmax函数处理logits，再取最大值索引：

# 转换为概率分布（dim=1表示每个样本的6个值之和为1）
probabilities = torch.nn.functional.softmax(output, dim=1)
# 获取概率最大的类别索引
preds_from_prob = probabilities.argmax(dim=1)
print(preds_from_prob)
# 输出：tensor([3, 3, 3, 3])

注意事项 #

你的标签中出现了值6，但输出仅为6列（对应类别索引0-5），需确认标签是否存在索引偏移，比如是否需要将标签减1才能与输出的类别索引对应。

内容的提问来源于stack exchange，提问作者prut