基于PyTorch的不平衡数据集图像分类实战指南

引言

在计算机视觉领域，图像分类任务常面临数据不平衡的挑战：某些类别的样本数量远多于其他类别（如医疗影像中正常样本占比90%，病变样本仅10%）。这种不平衡会导致模型偏向多数类，降低少数类的分类性能。PyTorch作为深度学习领域的核心框架，提供了灵活的工具链来应对此类问题。本文将从数据增强、损失函数设计、采样策略和模型架构优化四个维度，系统阐述基于PyTorch的不平衡数据集图像分类解决方案。

一、数据预处理与增强策略

1.1 类别权重分析

首先需量化数据不平衡程度。通过统计各类别样本数，计算类别频率分布：

import torch
from collections import Counter
# 假设labels为所有样本的类别标签列表
label_counts = Counter(labels)
class_weights = {cls: 1/count for cls, count in label_counts.items()}
normalized_weights = torch.tensor([class_weights[cls] for cls in labels], dtype=torch.float32)

此代码计算每个类别的逆频率权重，为后续加权损失函数提供基础。

1.2 针对性数据增强

对少数类样本实施更激进的数据增强：

from torchvision import transforms
# 基础变换（适用于所有样本）
base_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor()
])
# 少数类增强（仅对少数类样本应用）
augmented_transform = transforms.Compose([
    transforms.RandomRotation(15),
    transforms.ColorJitter(brightness=0.2, contrast=0.2),
    transforms.RandomAffine(degrees=0, translate=(0.1, 0.1))
])
# 实际应用时需根据标签动态选择变换
def get_transform(label):
    if label in minority_classes:  # 假设已定义minority_classes
        return transforms.Compose([base_transform, augmented_transform])
    return base_transform

这种差异化增强策略可有效扩充少数类样本的多样性。

二、损失函数优化

2.1 加权交叉熵损失

PyTorch的CrossEntropyLoss支持类别权重：

import torch.nn as nn
# 计算类别权重（需归一化）
class_counts = torch.tensor([label_counts[cls] for cls in range(num_classes)])
weights = 1. / class_counts.float()
weights = weights / weights.sum() * num_classes  # 归一化
# 创建加权损失函数
criterion = nn.CrossEntropyLoss(weight=weights)

该实现使少数类样本的损失贡献更大，平衡训练过程。

2.2 Focal Loss实现

针对极端不平衡场景，Focal Loss通过动态调节权重聚焦难样本：

class FocalLoss(nn.Module):
    def __init__(self, alpha=0.25, gamma=2.0, reduction='mean'):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.reduction = reduction
    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
        if self.reduction == 'mean':
            return focal_loss.mean()
        elif self.reduction == 'sum':
            return focal_loss.sum()
        return focal_loss

alpha控制类别平衡，gamma调节难样本聚焦程度，通常设为2.0。

三、采样策略实现

3.1 过采样与欠采样

PyTorch的WeightedRandomSampler可实现动态采样：

from torch.utils.data import WeightedRandomSampler
# 计算样本权重（少数类样本权重更高）
samples_weight = torch.tensor([class_weights[label] for label in labels])
sampler = WeightedRandomSampler(
    weights=samples_weight,
    num_samples=len(samples_weight),
    replacement=True  # 过采样需要replacement=True
)
# 创建DataLoader时使用sampler
dataloader = torch.utils.data.DataLoader(
    dataset,
    batch_size=32,
    sampler=sampler
)

此方法通过重复采样少数类样本平衡批次分布。

3.2 两阶段训练法

结合过采样与欠采样的混合策略：

# 第一阶段：过采样少数类
minority_indices = [i for i, label in enumerate(labels) if label in minority_classes]
majority_indices = [i for i, label in enumerate(labels) if label not in minority_classes]
# 复制少数类样本（假设复制比例为2倍）
augmented_indices = minority_indices * 2
balanced_indices = augmented_indices + majority_indices[:len(augmented_indices)]
# 第二阶段：欠采样多数类
sampled_majority = random.sample(majority_indices, len(minority_indices))
final_indices = minority_indices + sampled_majority

该方法先扩充少数类，再从多数类中随机采样等量样本，构建平衡数据集。

四、模型架构优化

4.1 类别平衡模块

在模型中嵌入类别注意力机制：

class ClassBalanceModule(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        self.class_weights = nn.Parameter(torch.ones(num_classes))
    def forward(self, x):
        # x的形状为[batch_size, num_classes]
        weighted_logits = x * self.class_weights.view(1, -1)
        return weighted_logits
# 在模型末尾使用
model = nn.Sequential(
    # ... 原有特征提取层 ...
    ClassBalanceModule(num_classes)
)

该模块通过可学习参数动态调整各类别的输出权重。

4.2 多任务学习框架

将分类任务与样本难度预测结合：

class MultiTaskModel(nn.Module):
    def __init__(self, backbone):
        super().__init__()
        self.backbone = backbone
        self.classifier = nn.Linear(512, num_classes)  # 假设特征维度为512
        self.difficulty_estimator = nn.Sequential(
            nn.Linear(512, 128),
            nn.ReLU(),
            nn.Linear(128, 1)
        )
    def forward(self, x):
        features = self.backbone(x)
        logits = self.classifier(features)
        difficulty = self.difficulty_estimator(features)
        return logits, difficulty

通过预测样本难度辅助分类器关注难样本，特别适用于不平衡场景。

五、评估指标选择

5.1 宏平均指标

优先使用宏平均F1分数而非准确率：

from sklearn.metrics import f1_score
# 计算宏平均F1
y_true = ...  # 真实标签
y_pred = ...  # 预测标签
macro_f1 = f1_score(y_true, y_pred, average='macro')

宏平均对每个类别同等加权，避免多数类主导评估结果。

5.2 混淆矩阵分析

可视化各类别分类情况：

import seaborn as sns
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix
cm = confusion_matrix(y_true, y_pred)
plt.figure(figsize=(10,8))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
plt.xlabel('Predicted')
plt.ylabel('True')
plt.show()

通过混淆矩阵可精准定位哪些少数类样本易被误分类。

六、完整训练流程示例

# 初始化模型
model = ResNet18(num_classes=10)  # 假设使用ResNet18
criterion = FocalLoss(alpha=0.25, gamma=2.0)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
# 创建平衡数据加载器
sampler = WeightedRandomSampler(...)  # 如前文所述
train_loader = DataLoader(dataset, batch_size=32, sampler=sampler)
# 训练循环
for epoch in range(100):
    model.train()
    for inputs, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
    # 验证阶段使用原始数据分布计算宏平均F1
    val_metrics = evaluate(model, val_loader)  # 自定义评估函数
    print(f"Epoch {epoch}: Val Macro F1 = {val_metrics['macro_f1']:.4f}")

七、实践建议

1. 渐进式平衡：从轻微过采样开始，逐步增加平衡强度，避免模型过拟合少数类
2. 早停机制：监控少数类的验证集性能作为早停依据
3. 集成学习：结合多个不平衡处理方法的模型集成通常效果更佳
4. 领域知识：利用数据先验知识设计更有效的增强策略（如医学图像中病变区域的局部增强）

结论

处理不平衡数据集的图像分类需要数据、算法、模型三方面的协同优化。PyTorch通过其灵活的张量操作和模块化设计，为实施加权损失、动态采样、类别平衡架构等高级技术提供了理想平台。实际应用中，建议采用”数据增强+损失函数调整+采样策略”的组合方案，并根据具体任务特点调整各组件参数。随着深度学习模型容量的提升，结合自监督预训练与微调策略将成为处理极端不平衡场景的新方向。