解决DeepCrack训练时RuntimeError：目标尺寸不匹配问题

错误核心原因 #

• 任务类型混淆：DeepCrack是语义分割模型（逐像素分类），但你当前用ImageFolder加载的是分类数据集，还错误地将输入图像当作标签传入损失函数，完全不符合分割任务逻辑。
• 模型输出处理错误：DeepCrack的forward方法返回(output, fuse5, fuse4, fuse3, fuse2, fuse1)，其中只有第一个output是最终分割结果，你用torch.stack(outputs, dim=0)把所有返回值堆叠，导致输出维度混乱，完全偏离模型预期输出格式。
• 损失函数与任务不匹配：你用了分类任务的CrossEntropyLoss，但分割任务需要逐像素的损失函数，两者输入输出维度逻辑完全不同。

分步修复方案 #

1. 修正模型输出提取 #

只保留DeepCrack返回的最终分割结果，丢弃中间融合层输出：

# 替换原输出处理代码
outputs = model(images)
final_output = outputs[0]  # 仅取最终分割输出，维度为[6,1,224,224]（二分类）

2. 重构分割数据集 #

分割任务需要图像-掩码对，而非分类数据集结构，你需要重新组织数据并自定义数据集类：

• 数据集结构示例：

  Data_Structure(Annotated)/
      images/
          img1.jpg
          img2.jpg
      masks/
          mask1.png  # 单通道掩码，像素值对应类别（0=背景，1=裂纹等）
          mask2.png
  

• 自定义Dataset类：

from PIL import Image

class CrackDataset(torch.utils.data.Dataset):
    def __init__(self, img_dir, mask_dir, transform=None):
        self.img_dir = img_dir
        self.mask_dir = mask_dir
        self.transform = transform
        self.img_names = os.listdir(img_dir)
    
    def __len__(self):
        return len(self.img_names)
    
    def __getitem__(self, idx):
        img_path = os.path.join(self.img_dir, self.img_names[idx])
        mask_path = os.path.join(self.mask_dir, self.img_names[idx].replace('.jpg', '.png'))
        
        image = Image.open(img_path).convert('RGB')
        mask = Image.open(mask_path).convert('L')  # 转为单通道灰度掩码
        
        if self.transform:
            image = self.transform(image)
            mask = transforms.ToTensor()(mask)  # 掩码无需归一化
        
        return image, mask

• 替换原数据集加载代码：

train_trans = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomVerticalFlip(p=0.5),
    transforms.RandomRotation(degrees=(0.5, 5)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])

# 加载分割数据集
train_dataset = CrackDataset(
    img_dir='../Data/Data_Structure(Annotated)/images',
    mask_dir='../Data/Data_Structure(Annotated)/masks',
    transform=train_trans
)

val_size = 127
train_size = len(train_dataset) - val_size
train_loader, val_loader = random_split(train_dataset, [train_size, val_size])
train_loader = torch.utils.data.DataLoader(train_loader, shuffle=True, batch_size=batch_size)
val_loader = torch.utils.data.DataLoader(val_loader, shuffle=True, batch_size=batch_size)

3. 匹配损失函数与任务类型 #

• 二分类任务（背景/裂纹）：使用BCEWithLogitsLoss，适配模型单通道输出：

criterion = nn.BCEWithLogitsLoss()

• 多分类任务（4种裂纹类型）：先修改DeepCrack的输出通道数，再用CrossEntropyLoss：

# 修改DeepCrack类的final层，输出通道改为4
self.final = Conv3X3(5, num_classes)  # num_classes=4

# 定义损失函数，训练时调整标签维度
criterion = nn.CrossEntropyLoss()

4. 修正训练循环逻辑 #

训练时使用掩码标签而非输入图像，同时匹配损失函数要求：

for epoch in range(1, n_epochs+1):
    train_loss = 0.0
    for data in train_loader:
        images, masks = data  # 取掩码作为标签
        optimizer.zero_grad()
        
        outputs = model(images)
        final_output = outputs[0]
        
        # 二分类场景
        loss = criterion(final_output, masks)
        # 多分类场景：
        # masks = masks.squeeze(1).long()  # 压缩通道维度并转为长整型
        # loss = criterion(final_output, masks)
        
        loss.backward()
        optimizer.step()
        train_loss += loss.item() * images.size(0)
    
    train_loss = train_loss/len(train_loader)
    print('Epoch: {} \tTraining Loss: {:.6f}'.format(epoch, train_loss))

验证输出维度 #

训练前先测试模型输出是否符合预期：

test_input = torch.randn(6, 3, 224, 224)
model = DeepCrack()
outputs = model(test_input)
print("最终输出维度:", outputs[0].shape)  # 二分类应为[6,1,224,224]，多分类应为[6,4,224,224]

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