二分类语义分割训练报错：BCELoss输入与标签尺寸不匹配

从报错信息可以明确看到，模型输出尺寸[4, 2, 1024, 1024]与标签尺寸[4, 1, 1024, 1024]不匹配，同时损失函数的选择和标签处理逻辑也存在问题，以下是两种可行的解决方案：

方案一：使用CrossEntropyLoss（推荐，符合二分类语义分割常规做法） #

核心修改点 #

1. 更换损失函数：CrossEntropyLoss适配模型输出的2通道类别概率，标签只需提供类别索引（0=背景，1=目标）
2. 修正标签处理逻辑：移除标签的通道维度，确保标签为整数类型的类别索引

具体代码修改 #

1. main.py

将BCELoss替换为CrossEntropyLoss：

# 原代码
# criterion = nn.BCELoss()
# 修改后
criterion = nn.CrossEntropyLoss()

2. dataset.py

• 导入numpy：在文件顶部添加import numpy as np
• 修改__getitem__方法，单独处理标签：

def __getitem__(self, idx):
    img_name = os.path.join(self.image_folder, self.images[idx])
    label_name = os.path.join(self.label_folder, self.labels[idx])

    image = Image.open(img_name).convert("RGB")
    # 确保标签为灰度图
    label = Image.open(label_name).convert("L")

    # 仅对图像应用transform
    if self.transform:
        image = self.transform(image)
    
    # 标签处理：转换为整数张量，移除通道维度，将255（白色目标）转为1
    label_np = np.array(label)
    label_np[label_np == 255] = 1
    label = torch.tensor(label_np, dtype=torch.long)  # 形状为[H, W]

    return image, label

3. train.py

无需修改损失计算逻辑，CrossEntropyLoss会自动匹配[B, C, H, W]的模型输出和[B, H, W]的标签。

方案二：使用BCELoss/BCEWithLogitsLoss（单通道输出） #

核心修改点 #

1. 修改模型输出通道：将输出通道改为1，对应二分类的概率输出
2. 调整标签维度：确保标签与模型输出维度一致（[B, 1, H, W]）
3. 匹配损失函数：用BCEWithLogitsLoss（无需手动加sigmoid，数值更稳定）或BCELoss（需加sigmoid激活）

具体代码修改 #

1. engine.py

修改模型输出通道为1，可选添加sigmoid（推荐用BCEWithLogitsLoss则不需要）：

class Model(nn.Module):
    def __init__(self, num_classes):
        super(Model, self).__init__()
        self.model = deeplabv3_resnet101(pretrained=True)
        # 将输出通道改为1
        self.model.classifier[-1] = nn.Conv2d(256, 1, kernel_size=(1, 1), stride=(1, 1))
        # 若用BCELoss则需要sigmoid，用BCEWithLogitsLoss则注释掉
        # self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        out = self.model(x)['out']
        # 若用BCELoss则返回self.sigmoid(out)，否则直接返回out
        return out

2. main.py

选择对应的损失函数：

# 用BCEWithLogitsLoss（无需sigmoid，推荐）
criterion = nn.BCEWithLogitsLoss()

# 若用BCELoss（需模型加sigmoid）
# criterion = nn.BCELoss()

3. dataset.py

• 导入numpy：在文件顶部添加import numpy as np
• 修改__getitem__方法，将标签转为[1, H, W]的浮点张量：

def __getitem__(self, idx):
    img_name = os.path.join(self.image_folder, self.images[idx])
    label_name = os.path.join(self.label_folder, self.labels[idx])

    image = Image.open(img_name).convert("RGB")
    label = Image.open(label_name).convert("L")

    if self.transform:
        image = self.transform(image)
    
    # 标签处理：归一化到0-1，添加通道维度
    label_np = np.array(label) / 255.0
    label = torch.tensor(label_np, dtype=torch.float32).unsqueeze(0)  # 形状为[1, H, W]

    return image, label

4. train.py

修改预测结果的计算逻辑（单通道输出用阈值判断类别）：

# 原代码
# preds = torch.argmax(outputs, dim=1)
# 修改后
preds = (outputs > 0.5).float().squeeze(1)  # 形状为[B, H, W]

额外注意事项 #

• 确保标签的像素值映射正确：如果原标签白色为255、黑色为0，必须转换为1（目标）和0（背景），否则损失计算会完全错误。
• 若使用数据增强（如Resize、Flip），需确保图像和标签执行完全相同的变换，可使用torchvision.transforms.v2中的方法同时处理图像和标签。

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