Softmax 输出层哪个维度？

问题描述

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

我使用的是softmax，所以输出层如下：

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

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

[batch_size,num_classes,width,height,depth]

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

谢谢！

解决方法

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