PyTorch如何实现让神经网络仅输出0、1、2三个整数？

你完全不需要搭建587个独立的三分类神经网络，用现有结构稍作调整就可以满足需求：

• 首先调整输出层配置：将你当前nn.Linear(hidden_size, num_classes)中的num_classes设置为587 * 3，前向传播时把输出结果reshape为(batch_size, 587, 3)的张量，这个张量最后一维的三个值对应587个位置各自的三分类logits。
• 训练阶段直接用nn.CrossEntropyLoss计算损失即可，不需要手动处理one-hot编码：只需要保证你的标签张量形状为(batch_size, 587)，每个位置的取值是0/1/2的真实整数标签，损失函数会自动处理分类计算逻辑。
• 推理阶段要得到0/1/2的整数输出，只需要对logits张量的最后一维取argmax(dim=-1)，得到的结果就是形状为(batch_size, 587)的整数向量，完全符合你的输出要求。

示例代码参考 #

# 网络结构定义修改
self.l1 = nn.Linear(input_size, hidden_size)
self.relu = nn.ReLU()
self.l2 = nn.Linear(hidden_size, 587 * 3) # 修改输出维度

# 前向传播逻辑
def forward(self, x):
    x = self.l1(x)
    x = self.relu(x)
    x = self.l2(x)
    return x.reshape(x.shape[0], 587, 3)

# 训练时损失计算
criterion = nn.CrossEntropyLoss()
logits = model(input_data)
# CrossEntropyLoss要求类别维度放在第二维，调整维度顺序即可
loss = criterion(logits.permute(0, 2, 1), true_labels)

# 推理时获取整数输出
pred_result = logits.argmax(dim=-1)

你提到的SoftMax输出概率问题只是因为你缺少了取最大值索引的步骤，上述方案已经覆盖了从训练到推理的全流程需求，和你设想的587个三分类网络效果等价，但是参数共享、计算效率远高于独立搭建多个网络。

内容的提问来源于stack exchange，提问作者FluidMechanics Potential Flows