python - Softmax 输出层。哪个维度？

Question

我有一个关于用于图像分割的神经元网络的问题。我正在使用可以在这里找到的 Deeplab 的 3D 实现

我正在使用softmax，所以输出层如下：

elif self.last_activation.lower() == 'softmax':
        output = nn.Softmax()(output)

没有定义维度，所以我想手动定义它。但我不确定我需要设置哪个维度。输出张量的维度如下：

[batch_size, num_classes, width, height, depth]

所以我认为这dim=1是正确的。那是对的吗？

谢谢！

标签： pythonneural-networkpytorchimage-segmentationdeeplab

Answer 1

实际上它应该是 1，因为您希望该轴总和为 1。
如果您需要使用交叉熵损失来训练您的网络，请小心，因为后者已经包含一个 softmax。