AI分类模型不平衡数据集进阶处理方法问询（TensorFlow/PyTorch）

我正在处理一个高度不平衡数据集的机器学习分类任务，担心模型性能表现。已经尝试过过采样、欠采样、类别加权等基础方法，现在想了解针对AI分类模型的更高级或高效的不平衡数据集处理方法。

以下是我用TensorFlow编写的Python代码片段：

import numpy as np
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense

# Load and preprocess the imbalanced dataset
# ...

# Define and compile the model
model = Sequential([
    Dense(64, activation='relu', input_shape=(input_dim,)),
    Dense(32, activation='relu'),
    Dense(num_classes, activation='softmax')
])

model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

# Train the model
model.fit(X_train, y_train, epochs=10, batch_size=32, verbose=1)

请问在TensorFlow或PyTorch等深度学习框架中，有哪些可提升模型性能与鲁棒性的针对性技术或进阶方案？

一、自定义损失函数：聚焦难分类样本 #

基础类别加权交叉熵对所有错误样本权重一致，Focal Loss通过降低易分类样本的权重，让模型优先关注少数类和难分类样本，是不平衡场景的核心优化手段。

TensorFlow实现：

class FocalLoss(tf.keras.losses.Loss):
    def __init__(self, alpha=0.25, gamma=2.0, name='focal_loss'):
        super().__init__(name=name)
        self.alpha = alpha
        self.gamma = gamma

    def call(self, y_true, y_pred):
        y_true = tf.one_hot(tf.cast(y_true, tf.int32), depth=y_pred.shape[-1])
        ce = tf.keras.losses.categorical_crossentropy(y_true, y_pred)
        p = tf.reduce_max(y_pred, axis=-1)
        pt = tf.where(tf.equal(y_true, 1), p, 1 - p)
        loss = self.alpha * tf.pow(1 - pt, self.gamma) * ce
        return tf.reduce_mean(loss)

# 替换原有损失函数
model.compile(optimizer='adam',
              loss=FocalLoss(alpha=0.25, gamma=2.0),
              metrics=['accuracy'])

PyTorch实现：

import torch
import torch.nn as nn

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, y_pred, y_true):
        ce_loss = nn.CrossEntropyLoss(reduction='none')(y_pred, y_true)
        pt = torch.exp(-ce_loss)
        focal_loss = self.alpha * (1 - pt)**self.gamma * ce_loss
        return torch.mean(focal_loss) if self.reduction == 'mean' else focal_loss

# 初始化损失函数
criterion = FocalLoss(alpha=0.25, gamma=2.0)

二、进阶采样策略：动态调整样本分布 #

1. 困难负样本挖掘

筛选模型分类错误的难样本（多数类中的难分类样本、所有少数类样本），优先用这些样本更新模型，提升对少数类的识别能力。

TensorFlow实现（结合tf.data）：

def hard_negative_mining(dataset, model, batch_size):
    def get_loss(sample):
        x, y = sample
        y_pred = model(x, training=False)
        loss = tf.keras.losses.sparse_categorical_crossentropy(y, y_pred)
        return loss, x, y

    loss_ds = dataset.map(get_loss)
    sorted_ds = loss_ds.shuffle(len(dataset)).sorted(lambda loss, _, __: -loss)
    hard_ds = sorted_ds.map(lambda _, x, y: (x, y)).batch(batch_size)
    return hard_ds

# 使用难样本数据集训练
hard_train_ds = hard_negative_mining(train_ds, model, 32)
model.fit(hard_train_ds, epochs=10)

PyTorch实现（自定义Sampler）：

from torch.utils.data import WeightedRandomSampler

def get_hard_sampler(dataloader, model):
    model.eval()
    losses = []
    with torch.no_grad():
        for x, y in dataloader:
            y_pred = model(x)
            loss = nn.CrossEntropyLoss(reduction='none')(y_pred, y)
            losses.extend(loss.tolist())
    weights = torch.tensor(losses)
    sampler = WeightedRandomSampler(weights, num_samples=len(weights), replacement=True)
    return sampler

# 生成难样本采样器并训练
hard_sampler = get_hard_sampler(train_loader, model)
hard_train_loader = torch.utils.data.DataLoader(train_dataset, sampler=hard_sampler, batch_size=32)

for epoch in range(10):
    model.train()
    for x, y in hard_train_loader:
        optimizer.zero_grad()
        y_pred = model(x)
        loss = criterion(y_pred, y)
        loss.backward()
        optimizer.step()

2. 自适应混合采样

训练初期对少数类过采样，后期逐步降低过采样比例；同时对多数类选择性欠采样，仅保留难分类样本，平衡训练过程中的样本分布。

三、模型结构优化：增强少数类特征捕捉 #

1. 添加注意力机制

让模型自动关注对分类关键的特征，尤其是少数类样本的独特特征。

TensorFlow实现：

class AttentionLayer(tf.keras.layers.Layer):
    def __init__(self, units):
        super().__init__()
        self.W = Dense(units)
        self.V = Dense(1)

    def call(self, inputs):
        inputs = tf.expand_dims(inputs, axis=1)
        score = self.V(tf.nn.tanh(self.W(inputs)))
        attention_weights = tf.nn.softmax(score, axis=1)
        context_vector = attention_weights * inputs
        return tf.reduce_sum(context_vector, axis=1)

# 修改原有模型
model = Sequential([
    Dense(64, activation='relu', input_shape=(input_dim,)),
    AttentionLayer(32),
    Dense(32, activation='relu'),
    Dense(num_classes, activation='softmax')
])

PyTorch实现：

class AttentionLayer(nn.Module):
    def __init__(self, units):
        super().__init__()
        self.W = nn.Linear(units, units)
        self.V = nn.Linear(units, 1)

    def forward(self, x):
        x = x.unsqueeze(1)
        score = torch.tanh(self.W(x))
        attention_weights = torch.softmax(self.V(score), dim=1)
        context_vector = attention_weights * x
        return torch.sum(context_vector, dim=1)

# 修改原有模型
model = nn.Sequential(
    nn.Linear(input_dim, 64),
    nn.ReLU(),
    AttentionLayer(32),
    nn.Linear(32, 32),
    nn.ReLU(),
    nn.Linear(32, num_classes),
    nn.Softmax(dim=1)
)

2. 对比学习预训练

先通过自监督对比学习让模型学习通用特征，再用少量标签数据微调分类器，有效利用未标注数据，提升对少数类的特征提取能力。

TensorFlow实现（简化版）：

# 定义对比学习模型
class ContrastiveModel(tf.keras.Model):
    def __init__(self, feature_dim=64):
        super().__init__()
        self.encoder = Sequential([
            Dense(128, activation='relu', input_shape=(input_dim,)),
            Dense(feature_dim, activation='relu')
        ])
        self.projection_head = Sequential([Dense(64, activation='relu'), Dense(32)])

    def call(self, x):
        return self.projection_head(self.encoder(x))

# NT-Xent对比损失
def nt_xent_loss(z1, z2, temperature=0.1):
    z = tf.concat([z1, z2], axis=0)
    cos_sim = tf.matmul(z, tf.transpose(z)) / temperature
    mask = tf.eye(tf.shape(z)[0], dtype=tf.bool)
    cos_sim = tf.where(mask, -1e9 * tf.ones_like(cos_sim), cos_sim)
    labels = tf.concat([tf.range(tf.shape(z1)[0]), tf.range(tf.shape(z1)[0])], axis=0)
    return tf.reduce_mean(tf.keras.losses.sparse_categorical_crossentropy(labels, cos_sim, from_logits=True))

# 预训练对比模型
contrast_model = ContrastiveModel()
optimizer = tf.keras.optimizers.Adam()

@tf.function
def train_step(x1, x2):
    with tf.GradientTape() as tape:
        loss = nt_xent_loss(contrast_model(x1), contrast_model(x2))
    gradients = tape.gradient(loss, contrast_model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, contrast_model.trainable_variables))
    return loss

# x1、x2为同一样本的两种增强版本（如加噪声、裁剪）
for epoch in range(50):
    loss = train_step(x1_augmented, x2_augmented)

# 用预训练编码器初始化分类模型
classifier_model = Sequential([
    contrast_model.encoder,
    Dense(32, activation='relu'),
    Dense(num_classes, activation='softmax')
])
classifier_model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
classifier_model.fit(X_train, y_train, epochs=10)

四、训练策略调整：提升模型鲁棒性 #

1. 标签平滑

给标签添加小噪声，降低模型对多数类的过度自信，减少过拟合，提升对少数类的泛化能力。

TensorFlow实现：

model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(label_smoothing=0.1),
              metrics=['accuracy'])

PyTorch实现：

criterion = nn.CrossEntropyLoss(label_smoothing=0.1)

2. 混合数据增强（MixUp/CutMix）

混合不同样本的特征与标签生成新样本，增加数据多样性，帮助模型学习通用特征，减少对多数类的依赖。

TensorFlow实现（MixUp）：

def mixup_data(x, y, alpha=0.2):
    lam = np.random.beta(alpha, alpha)
    batch_size = tf.shape(x)[0]
    index = tf.random.shuffle(tf.range(batch_size))
    mixed_x = lam * x + (1 - lam) * tf.gather(x, index)
    mixed_y = lam * tf.one_hot(y, num_classes) + (1 - lam) * tf.one_hot(tf.gather(y, index), num_classes)
    return mixed_x, mixed_y

@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        x_mixed, y_mixed = mixup_data(x, y)
        y_pred = model(x_mixed, training=True)
        loss = tf.reduce_mean(tf.keras.losses.categorical_crossentropy(y_mixed, y_pred))
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

# 训练循环
for epoch in range(10):
    for x_batch, y_batch in train_ds:
        loss = train_step(x_batch, y_batch)

五、评估指标优化：避免被准确率误导 #

不平衡数据集下准确率无参考价值，需使用以下指标：

• 加权F1-Score：综合精确率与召回率
• Precision-Recall曲线：更适合极不平衡场景
• 混淆矩阵：直观展示类别分类情况
• Cohen's Kappa：衡量模型与随机分类的差异

TensorFlow/PyTorch通用指标计算：

from sklearn.metrics import f1_score, confusion_matrix, cohen_kappa_score

# TensorFlow：获取预测结果
y_pred = model.predict(X_test)
y_pred_classes = tf.argmax(y_pred, axis=1).numpy()

# PyTorch：获取预测结果
# model.eval()
# y_pred_classes = []
# with torch.no_grad():
#     for x, y in test_loader:
#         y_pred_classes.extend(torch.argmax(model(x), dim=1).tolist())

# 计算指标
f1 = f1_score(y_test, y_pred_classes, average='weighted')
cm = confusion_matrix(y_test, y_pred_classes)
kappa = cohen_kappa_score(y_test, y_pred_classes)

print(f"Weighted F1-Score: {f1:.4f}")
print(f"Cohen's Kappa: {kappa:.4f}")
print("Confusion Matrix:\n", cm)

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