使用smp.Unet配置aux_params时出现AttributeError: 'tuple'无size属性

错误现象 #

使用segmentation_models_pytorch（简称smp）的Unet模型并传入aux_params参数后，训练时出现如下错误：

File .../python3.11/site-packages/segmentation_models_pytorch/utils/train.py:51, in Epoch.run(self, dataloader)
     49 for x, y in iterator:
     50     x, y = x.to(self.device), y.to(self.device)
---> 51     loss, y_pred = self.batch_update(x, y)
     53     # update loss logs
     54     loss_value = loss.cpu().detach().numpy()
...
-> 3162 if not (target.size() == input.size()):
   3163     raise ValueError("Target size ({}) must be the same as input size ({})".format(target.size(), input.size()))
   3165 return torch.binary_cross_entropy_with_logits(input, target, weight, pos_weight, reduction_enum)
File ".../train_model.py", line 153, in train
    train_logs = self.train_epoch.run(self.train_loader)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../train_model.py", line 173, in main
    water_seg_model.train(epoch_number=100)
  File ".../train_model.py", line 176, in <module>
    main()
AttributeError: 'tuple' object has no attribute 'size'

错误原因 #

当给smp.Unet传入aux_params时，模型的输出不再是单个张量，而是元组(主分割输出, 辅助分类输出)。但默认的smp.utils.train.TrainEpoch和损失函数只会处理单个输出张量，直接将元组传入损失函数时，损失函数会尝试调用tuple.size()方法，而元组没有这个属性，因此触发错误。

解决方法 #

需要自定义损失函数来处理模型的双输出，同时调整指标计算逻辑只针对主分割输出：

1. 自定义组合损失函数 #

编写损失函数，拆分模型输出的元组，分别计算主损失和辅助损失后加权求和：

def combined_loss(y_pred, y_true):
    # 拆分模型输出：主分割输出、辅助分支输出
    main_output, aux_output = y_pred
    # 计算主分割任务损失
    main_loss = smp.losses.SoftBCEWithLogitsLoss()(main_output, y_true)
    # 计算辅助分支损失（若辅助分支是全局分类，需调整目标格式适配输出尺寸）
    aux_loss = smp.losses.SoftBCEWithLogitsLoss()(aux_output, y_true)
    # 加权求和，辅助损失权重可根据需求调整
    return main_loss + 0.3 * aux_loss

# 给损失函数命名，方便TrainEpoch记录训练日志
combined_loss.__name__ = 'CombinedLoss'

2. 自定义指标处理逻辑 #

默认指标会接收模型的元组输出，需要修改指标只使用主分割输出计算：

class CustomIoU(smp.utils.metrics.IoU):
    def __call__(self, y_pred, y_true):
        # 只取主分割输出计算IoU指标
        main_output, _ = y_pred
        return super().__call__(main_output, y_true)

3. 更新模型初始化代码 #

将自定义的损失和指标替换原有的配置：

class SegmentationModel():
    def __init__(self):
        # ... 原有模型初始化代码保持不变 ...
        
        # 替换为自定义组合损失
        self.loss = combined_loss
        
        # 替换为自定义IoU指标
        self.metrics = [
            CustomIoU(threshold=THRESHOLD),
        ]
        
        # ... 其余代码保持不变 ...

补充说明 #

如果辅助分支是全局分类任务（而非与主输出同尺寸的分割），还需要修改数据集返回的目标，添加全局分类标签，或者调整辅助损失的计算逻辑适配输出尺寸。

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