pytorch - 图像分类器模型的输出大小应该是多少？

Question

我正在执行图像分类任务。图像被标记为 0 1 2。模型输出中最后一个线性层的大小应该是 3 还是 1 ？一般来说，对于 3 类操作，输出设置为 3，作为这三个的结果，返回最大概率。但是我看到在某些代码中最后一层设置为1。我认为这实际上是合乎逻辑的。你有什么想法 ？（我也不在最后一层使用 softmax 或 sigmoid 函数。）

Answer 1

要对类进行分类c（c在您的示例中为 = 3），您需要预测每个类的概率，因此您需要输出c-dim 输出。
通常你不会明确地应用softmax到“原始预测”（又名“logits”）——损失函数通常会以一种在数值上更稳健的方式为你做这件事（参见，例如，nn.CrossEntropyLoss）。

训练模型后，在推理时，您可以接管argmax预测的clogits 并输出单个标量 - 预测类的索引。这只能在推理期间完成，因为argmax它不是可微分的操作。