嘿，针对你遇到的类别不平衡问题，自定义加权损失确实是个非常有效的解决思路。我来一步步帮你把这个功能集成到现有代码里，跟着改就行：

第一步：计算类别权重 #

首先我们需要根据训练集的类别分布计算每个类的权重，核心思路是给样本量少的类别更高的权重。这里有两种简单的实现方式：

方式1：手动计算逆频率权重 #

from collections import Counter

# 先把编码后的标签转回类别索引（如果你的y_train是one-hot格式）
train_label_indices = y_train.argmax(dim=1).numpy()
# 统计每个类别的样本数
label_counts = Counter(train_label_indices)
total_samples = len(train_label_indices)
# 计算权重：总样本数/类别样本数
class_weights = torch.tensor(
    [total_samples / label_counts[idx] for idx in sorted(label_counts.keys())],
    dtype=torch.float32
).to(cuda)  # 移到GPU上和模型保持一致

方式2：用sklearn工具自动计算 #

sklearn提供了现成的compute_class_weight函数，能自动生成平衡权重：

from sklearn.utils.class_weight import compute_class_weight

# 获取训练集的类别索引
train_label_indices = y_train.argmax(dim=1).numpy()
# 计算平衡权重
class_weights = compute_class_weight(
    class_weight='balanced',
    classes=sorted(set(train_label_indices)),
    y=train_label_indices
)
# 转成PyTorch张量并移到GPU
class_weights = torch.tensor(class_weights, dtype=torch.float32).to(cuda)

第二步：定义加权损失函数 #

你有两种选择，要么直接用PyTorch自带的工具，要么完全自定义实现：

选择1：用PyTorch自带的加权交叉熵损失 #

PyTorch的CrossEntropyLoss本身就支持传入类别权重，这是最省心的方式：

# 初始化加权损失函数
weighted_loss_fn = torch.nn.CrossEntropyLoss(weight=class_weights)

选择2：自定义基于logits的损失函数 #

如果你想手动实现（比如更灵活地调整逻辑），可以自己写一个函数：

def weighted_cross_entropy(logits, labels, class_weights):
    # 如果标签是one-hot格式，先转成类别索引
    if labels.dim() == 2:
        labels = labels.argmax(dim=1).to(cuda)
    # 计算log softmax
    log_probs = torch.nn.functional.log_softmax(logits, dim=1)
    # 取出对应类别的log概率值
    selected_log_probs = log_probs.gather(1, labels.unsqueeze(1)).squeeze(1)
    # 乘以对应类别的权重
    weighted_log_probs = selected_log_probs * class_weights[labels]
    # 计算最终损失（取负后求均值）
    loss = -weighted_log_probs.mean()
    return loss

第三步：修改训练循环，替换默认损失计算 #

原来的代码中，调用bert()时传入labels会自动计算默认损失，现在我们要替换成自己的加权损失，修改后的训练循环如下：

for epoch in range(NUM_EPOCHS):
    epoch_losses = []
    for x, y in tqdm(_load_data(x_train, y_train, batch_size=10)):
        bert.zero_grad()
        # 这里不再传入labels参数，只获取模型输出的logits
        out = bert(x, attention_mask=x.ne(tokenizer.pad_token_id).to(int))
        logits = out.logits
        
        # --- 二选一：使用自定义损失 ---
        # 方式1：用自带的加权损失函数
        y_label_indices = y.argmax(dim=1).to(cuda)  # 转成类别索引并移到GPU
        loss = weighted_loss_fn(logits, y_label_indices)
        
        # 方式2：用自己写的损失函数
        # loss = weighted_cross_entropy(logits, y, class_weights)
        
        epoch_losses.append(loss.item())
        loss.backward()
        optim.step()
    # 打印每个epoch的平均损失，方便监控训练情况
    print(f"Epoch {epoch+1} 训练完成，平均损失：{mean(epoch_losses):.4f}")

几个需要注意的细节 #

• 确保所有张量都在同一设备上：模型在GPU（cuda）上，所以标签张量也要用.to(cuda)移到GPU，否则会报错。
• 标签格式要对应：CrossEntropyLoss接受的是类别索引（一维张量），如果你的y_train是one-hot格式（二维张量），一定要转成索引格式。
• 验证权重合理性：可以打印class_weights看看，样本量少的类别权重应该明显高于样本多的类别，比如760样本的类权重接近1，10样本的类权重应该是76左右。

备注：内容来源于stack exchange，提问作者Hoang Cuong Nguyen