PyTorch实现特定样本加权的自定义BCEWithLogitsLoss方法

核心思路 #

将「文件名含@的待惩罚样本」标识整合进PyTorch标准数据流：

• 改造数据集类，每个batch随图像、标签同步返回样本惩罚标识
• 自定义损失支持逐样本计算，对满足条件的样本施加lambda权重惩罚
• 训练循环仅需调整batch解包逻辑，完全兼容常规训练范式

分步实现 #

1. 扩展数据集，返回样本惩罚标记 #

自定义Dataset的__getitem__方法，在遍历文件时提前标记含@的样本，返回值新增0/1惩罚标识位：

from torch.utils.data import Dataset
from PIL import Image
import os
import torch

class FaceDetDataset(Dataset):
    def __init__(self, data_root, transform=None):
        self.data_root = data_root
        self.transform = transform
        self.sample_list = []

        # 按自己的目录结构遍历样本，以下为二分类示例：0=非人脸，1=人脸
        for cls_label in [0, 1]:
            cls_path = os.path.join(data_root, str(cls_label))
            for fname in os.listdir(cls_path):
                fpath = os.path.join(cls_path, fname)
                # 文件名含@则标记为待惩罚样本
                penalize_flag = 1 if '@' in fname else 0
                self.sample_list.append((fpath, cls_label, penalize_flag))

    def __len__(self):
        return len(self.sample_list)

    def __getitem__(self, index):
        fpath, label, penalize_flag = self.sample_list[index]
        img = Image.open(fpath).convert('RGB')
        if self.transform:
            img = self.transform(img)
        return img, torch.tensor(label, dtype=torch.float32), torch.tensor(penalize_flag, dtype=torch.float32)

2. 实现自定义加权损失函数 #

修改原生BCEWithLogitsLoss逻辑：先关闭默认聚合逐样本计算损失，再根据规则给对应样本加权，最后做归一化聚合。如果需要实现仅预测错误时才对@样本施加惩罚，可以在无梯度环境下计算错分mask，和惩罚标识取交集后再应用权重：

import torch.nn as nn
import torch.nn.functional as F
from typing import Optional

class PenalizedBCEWithLogitsLoss(nn.Module):
    def __init__(self, penalty_lambda: float, pos_weight: Optional[torch.Tensor] = None, reduction: str = 'mean', penalize_wrong_only: bool = True):
        super().__init__()
        self.penalty_lambda = penalty_lambda
        self.pos_weight = pos_weight
        self.reduction = reduction
        self.penalize_wrong_only = penalize_wrong_only # 控制是否仅错分才加惩罚

    def forward(self, pred: torch.Tensor, target: torch.Tensor, penalize_mask: torch.Tensor) -> torch.Tensor:
        # 逐样本计算BCE损失，不做聚合，保留每个样本的独立损失值
        per_sample_loss = F.binary_cross_entropy_with_logits(
            pred.flatten(),
            target.flatten(),
            pos_weight=self.pos_weight,
            reduction='none'
        )

        # 构造逐样本权重矩阵
        sample_weights = torch.ones_like(per_sample_loss)
        penalize_mask = penalize_mask.flatten()

        if self.penalize_wrong_only:
            # 仅对预测错误的@样本施加惩罚
            with torch.no_grad():
                pred_label = (torch.sigmoid(pred.flatten()) > 0.5).float()
                wrong_pred_mask = (pred_label != target.flatten()).float()
                apply_penalty = penalize_mask * wrong_pred_mask
        else:
            # 所有@样本不管预测对错都加权
            apply_penalty = penalize_mask

        sample_weights[apply_penalty == 1] = self.penalty_lambda
        weighted_loss = per_sample_loss * sample_weights

        # 按指定规则聚合损失
        if self.reduction == 'mean':
            # 除以总权重和而非样本数，避免权重缩放改变梯度整体量级，和原生损失尺度对齐
            return weighted_loss.sum() / sample_weights.sum()
        elif self.reduction == 'sum':
            return weighted_loss.sum()
        else:
            return weighted_loss

3. 训练循环适配 #

训练时仅需在解包batch时新增接收惩罚标识，其余逻辑（梯度清零、反向传播、参数更新）和官方标准训练流程完全一致：

from torch.utils.data import DataLoader

# 初始化组件
train_dataset = FaceDetDataset(data_root='./train', transform=train_transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
model = get_pretrained_face_model() # 加载预训练模型、冻结指定层的逻辑按自己需求写
criterion = PenalizedBCEWithLogitsLoss(penalty_lambda=6.0) # lambda按实验需求调整，如6代表错分@样本惩罚6倍
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-4)

# 标准训练循环
model.train()
for epoch in range(total_epochs):
    running_loss = 0.0
    for imgs, labels, penalize_mask in train_loader:
        imgs = imgs.cuda()
        labels = labels.cuda()
        penalize_mask = penalize_mask.cuda()

        optimizer.zero_grad()
        preds = model(imgs).squeeze()
        loss = criterion(preds, labels, penalize_mask)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
    # 验证、日志打印、模型保存等逻辑和常规训练无差异

注意事项 #

• 该实现完全保留了原生BCEWithLogitsLoss的数值稳定性，不存在sigmoid溢出风险
• 逐样本加权逻辑不影响模型层的冻结/微调设置，可直接搭配预训练模型微调流程使用
• 验证/测试阶段不需要计算损失时，可直接忽略Dataset返回的惩罚标识，不影响推理逻辑

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