处理不同形状批次的损失计算：加权平均损失方法

import torch

# 模拟不同批次的损失
losses_perbatch = [torch.randn(8, 1), torch.randn(4, 1), torch.randn(2, 1)]

# 计算总样本数量
total_samples = sum([len(batch) for batch in losses_perbatch])

# 计算每个批次的加权平均损失
weighted_mean_perbatch = torch.tensor([batch.sum() for batch in losses_perbatch]) / total_samples

# 等价于：
# weighted_mean_perbatch = torch.tensor([batch.mean() * len(batch) for batch in losses_perbatch]) / total_samples

# 计算最终的加权平均损失
final_weighted_loss = sum(weighted_mean_perbatch)

print(f"最终加权平均损失: {final_weighted_loss}")

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)

        # Zero Gradients
        optimizer.zero_grad()

        # Forward pass
        y_first, y_second = model(x)

        losses = []
        batch_sizes = []  # 记录每个batch的有效样本数量

        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

            # Padding of -1 is removed from y
            indices_mask = y[j].ne(-1)
            valid_indices = y[j][indices_mask]

            valid_predicted = predicted[valid_mask]
            valid_actual = actual[valid_mask]

            loss = mae_fn(valid_predicted, valid_actual, valid_indices)

            losses.append(loss.sum()) # 存储loss的总和
            batch_sizes.append(len(valid_indices)) # 存储有效样本的数量

        # Backward pass and update
        total_samples_in_batch = sum(batch_sizes)
        weighted_losses = [loss / total_samples_in_batch * batch_size for loss, batch_size in zip(losses, batch_sizes)]
        loss = sum(weighted_losses)
        loss.backward()

        optimizer.step()

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

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

    # Return the average MAEs for y type
    return (
        np.mean(train_mae)
    )