PyTorch灰度图像像素级分类中Softmax输出异常的解决方法

核心问题：Softmax的维度参数错误 #

你的代码里F.softmax(x, 2)指定的维度完全不对，这直接导致Softmax没有作用在通道维度上，也就无法实现同一像素的多通道值求和为1的预期。

在PyTorch中，你的网络输出张量shape是(batch_size, channels, height, width)：

• 第0维：批量大小
• 第1维：通道数（对应你的2个分类标签）
• 第2维：图像高度
• 第3维：图像宽度

你需要对每个像素的通道维度做Softmax，也就是指定dim=1，这样每个像素的所有通道值才会被归一化，求和等于1。而你现在用dim=2，是对每个batch、每个通道的高度维度做Softmax，完全偏离了你的需求。

修正步骤 #

1. 修复Softmax维度
   把forward函数里的Softmax调用改成：

x = F.softmax(x, dim=1)

修改后，每个像素的两个通道值就会被正确归一化，求和接近1（浮点精度误差会导致微小偏差）。

2. 优化损失函数（强烈推荐）
   你当前用RMSE损失处理分类任务并不合适——RMSE是回归任务的损失函数，像素级分类更适合用交叉熵损失：

• 如果你想保留网络最后的Softmax，可以改用nn.NLLLoss（需要先把F.softmax换成F.log_softmax）
• 更高效的方式是直接去掉网络中的Softmax，使用nn.CrossEntropyLoss，它内部会自动计算LogSoftmax和负对数似然损失，数值稳定性更好。

示例代码如下：

# 修改网络forward函数，移除Softmax
def forward(self, x):
    in_size = x.size(0)
    x = F.relu(self.conv11(x))
    x = F.relu(self.conv12(x))
    x = self.conv13(x)  # 直接输出原始logits，不做归一化
    return x

# 定义损失函数
criterion = nn.CrossEntropyLoss()
# 训练时计算损失
loss = criterion(outputs, labels)

另外，针对你的二分类场景（仅0和1两个标签），还可以简化网络结构：把最后一层conv13改成单通道输出，配合nn.BCEWithLogitsLoss，这样更简洁高效：

# 修改网络最后一层为1个通道
self.conv13 = nn.Conv2d(256, 1, kernel_size=3, padding=1)

# forward函数
def forward(self, x):
    x = F.relu(self.conv11(x))
    x = F.relu(self.conv12(x))
    x = self.conv13(x)
    return x

# 定义损失函数
criterion = nn.BCEWithLogitsLoss()
# 注意labels需要转换为float类型并匹配输出shape
loss = criterion(outputs, labels.float().unsqueeze(1))

验证修正效果 #

完成上述修改后，再运行你的验证代码，同一像素的通道值之和就会符合预期（接近1）。

内容的提问来源于stack exchange，提问作者Farshid Rayhan