批量样本形状不同时如何计算损失？

问题核心在于每个样本的有效损失元素数量不同，直接堆叠会报错，而单样本取均值再平均会因样本权重失衡导致泛化差。正确的做法是收集所有样本的有效损失元素，再计算全局均值，让每个损失元素的权重一致。

修改方案 #

1. 调整训练函数中的损失收集逻辑 #

将原来收集每个样本的损失张量，改为收集所有样本的有效损失元素，最后统一计算全局均值：

def training():
    model.train()
    
    train_mae = []
    
    progress = tqdm(train_dataloader, desc='Training')
    for batch_index, batch in enumerate(progress):
        x = batch['x'].to(device)
        x_lengths = batch['x_lengths'].to(device)
        y = batch['y'].to(device)
        y_type = batch['y_type'].to(device)
        y_valid_indices = batch['y_valid_indices'].to(device)

        optimizer.zero_grad()

        y_first, y_second = model(x)

        # 收集所有有效损失元素，而非每个样本的损失张量
        all_loss_elements = []

        for j in range(len(x_lengths)):
            x_length = x_lengths[j].item()

            if y_type[j].item() == 0:
                predicted = y_first[j]
            else:
                predicted = y_second[j]

            actual = y[j]
            
            valid_mask = torch.zeros_like(predicted, dtype=torch.bool)
            valid_mask[:x_length] = 1
            
            indices_mask = y[j].ne(-1)
            valid_indices = y[j][indices_mask]

            valid_predicted = predicted[valid_mask]
            valid_actual = actual[valid_mask]
            
            # 获取当前样本的有效损失元素
            loss_elements = mae_fn(valid_predicted, valid_actual, valid_indices)
            all_loss_elements.append(loss_elements)

        # 拼接所有损失元素并计算全局均值
        loss = torch.cat(all_loss_elements).mean()
        loss.backward()

        optimizer.step()
        
        train_mae.append(loss.detach().cpu().numpy())

        progress.set_description(
            f"mae: {loss.detach().cpu().numpy():.4f}"
        )

    return np.mean(train_mae)

2. 保留原mae_fn不变 #

def mae_fn(output, target, indices):
    clipped_target = torch.clip(target, min=0, max=1)
    maes = F.l1_loss(output, clipped_target, reduction='none')
    return maes[indices]

为什么原来的方法会导致测试损失差？ #

原来对每个样本的损失先取均值再整体平均，相当于给有效数据点少的样本赋予了和数据点多的样本相同的权重。比如一个样本只有1个有效点，另一个有100个有效点，两者的平均损失在全局计算中占比相同，但实际上100个点的样本包含更多的任务信息，这种权重失衡会导致模型偏向拟合数据点少的样本，最终泛化能力下降，测试损失升高。

而使用全局均值的方式，每个有效损失元素的权重一致，模型能学习到更全面的任务信息，泛化效果更稳定。

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