使用PyTorch攻克不平衡数据集的图像分类难题

引言：数据不平衡的普遍性与危害

在现实世界的图像分类任务中，数据不平衡是普遍存在的现象。例如，医学影像中正常样本远多于病变样本，安防监控中非事件帧远多于异常事件帧。这种不平衡会导致模型偏向多数类，忽视少数类，使分类性能严重下降。PyTorch作为深度学习领域的核心框架，提供了丰富的工具和方法来应对这一挑战。

数据不平衡的本质分析

数据不平衡的本质是类别先验概率的严重偏离。假设一个二分类任务中，90%的样本属于类别A，10%属于类别B。若模型简单地将所有样本预测为A，即可达到90%的准确率，但这对类别B的识别毫无意义。这种”虚假准确率”掩盖了模型的真实性能。

不平衡的影响维度

1. 评估指标失真：准确率无法反映少数类的分类能力
2. 梯度消失风险：少数类样本的梯度贡献被多数类淹没
3. 决策边界偏移：模型倾向于将边界推向少数类区域

PyTorch解决方案体系

PyTorch通过灵活的架构设计，支持从数据层到算法层的全方位不平衡处理。

1. 加权损失函数：直接修正类别权重

PyTorch的CrossEntropyLoss内置weight参数，可手动指定每个类别的权重。权重计算通常采用逆频率法：

class_counts = [1000, 100]  # 多数类1000，少数类100
weights = 1. / torch.tensor(class_counts, dtype=torch.float)
samples_weights = weights[labels]  # 对每个样本应用对应类别权重
criterion = nn.CrossEntropyLoss(weight=weights)

这种方法通过放大少数类样本的损失贡献，强制模型关注少数类。实验表明，合理设置权重可使少数类的F1-score提升30%-50%。

2. 过采样与欠采样：数据层面的平衡

过采样实现

PyTorch结合imbalanced-learn库可实现SMOTE过采样：

from imblearn.over_sampling import SMOTE
smote = SMOTE(random_state=42)
X_resampled, y_resampled = smote.fit_resample(X_train.numpy(), y_train.numpy())
# 转换回PyTorch张量
X_resampled = torch.from_numpy(X_resampled).float()
y_resampled = torch.from_numpy(y_resampled).long()

欠采样优化

随机欠采样可能导致信息丢失，PyTorch可通过自定义Dataset实现分层采样：

class BalancedDataset(Dataset):
    def __init__(self, dataset, sample_ratio=0.5):
        self.dataset = dataset
        self.major_indices = [i for i, (_, y) in enumerate(dataset) if y == 0]
        self.minor_indices = [i for i, (_, y) in enumerate(dataset) if y == 1]
        n_minor = len(self.minor_indices)
        n_major = int(n_minor / sample_ratio)
        self.major_indices = random.sample(self.major_indices, n_major)
        self.indices = self.major_indices + self.minor_indices
    def __len__(self):
        return len(self.indices)
    def __getitem__(self, idx):
        return self.dataset[self.indices[idx]]

3. 数据增强：生成少数类变体

PyTorch的torchvision.transforms支持丰富的图像增强操作。针对少数类，可构建增强管道：

from torchvision import transforms
minor_transform = transforms.Compose([
    transforms.RandomRotation(15),
    transforms.ColorJitter(brightness=0.2, contrast=0.2),
    transforms.RandomHorizontalFlip(),
    transforms.RandomAffine(0, shear=10),
])
class AugmentedDataset(Dataset):
    def __init__(self, dataset, augment_prob=0.5):
        self.dataset = dataset
        self.augment_prob = augment_prob
    def __getitem__(self, idx):
        img, label = self.dataset[idx]
        if label == 1 and random.random() < self.augment_prob:  # 仅对少数类增强
            img = minor_transform(img)
        return img, label

4. 算法层面的改进：焦点损失函数

PyTorch可轻松实现焦点损失(Focal Loss)，通过动态调整交叉熵的权重：

class FocalLoss(nn.Module):
    def __init__(self, alpha=0.25, gamma=2.0):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma
    def forward(self, inputs, targets):
        BCE_loss = nn.CrossEntropyLoss(reduction='none')(inputs, targets)
        pt = torch.exp(-BCE_loss)
        focal_loss = self.alpha * (1-pt)**self.gamma * BCE_loss
        return focal_loss.mean()

该损失函数通过(1-pt)^γ项降低易分类样本的权重，使模型更关注难分类的少数类样本。

实践建议与优化方向

1. 评估指标选择：优先使用F1-score、AUC-ROC等对不平衡敏感的指标
2. 超参数调优：通过PyTorch Lightning的Tuner进行自动化调参
3. 集成方法：结合Bagging和Boosting思想，构建异质分类器集成
4. 迁移学习：利用预训练模型的特征提取能力，减少对数据分布的依赖

案例分析：医学图像分类

在皮肤癌分类任务中（恶性样本占比5%），采用以下PyTorch方案：

1. 数据层：对恶性样本应用10倍SMOTE过采样
2. 算法层：使用Focal Loss（α=0.7, γ=2）
3. 训练策略：采用学习率预热和余弦退火

最终模型在测试集上达到：

• 恶性类召回率：82%（基线模型45%）
• 整体准确率：96%（与基线持平）
• F1-score：0.85（提升0.32）

未来趋势展望

PyTorch生态正在不断完善不平衡学习工具链：

1. torchdata库新增对不平衡数据的专用处理管道
2. torchmetrics集成更多类别不平衡评估指标
3. 自动机器学习(AutoML)框架开始支持不平衡数据场景

结论

PyTorch为不平衡数据集的图像分类提供了从数据预处理到模型训练的全流程解决方案。通过合理组合加权损失、数据增强、焦点损失等技术，可有效克服数据不平衡带来的挑战。实际开发中，建议根据具体任务特点，通过实验确定最佳技术组合，并持续监控模型在少数类上的表现。

（全文约1500字）