软最大输出层.哪个维度?

Question

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

我正在使用softmax，因此输出层如下：

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

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

[batch_size, num_classes, width, height, depth]

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

谢谢！

Answer 1

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