PyTorch模型在首个数据集批次后无法学习新标签问题求助

我编写了一段代码，遍历按标签分组的不同数据集，让模型逐个学习每个标签对应的数据集。目前遇到的问题是：模型能完美学习首个数据集，但后续对新标签几乎没有学习效果（甚至完全没效果）。

核心代码 #

for index, datasetbatch in enumerate(dataset_batches): # 遍历数据集批次，第一个批次只有标签'0'的数据，后续依次是'1'、'2'等
    trainModel(params, model, datasetbatch[1], index + 1, device)

def trainModel(params: Parameters, model: CreatedModel, train_loader, currentBatchNum, device):
    criterion = nn.CrossEntropyLoss()
    optimizer = get_optimizer(params, model=model)

    fitnesses_each_epoch = []

    for i in range(params.epochs):
        fitnesses_in_epoch = []
        start_time = time.time()

        for b, (X_train, y_train) in enumerate(train_loader):
            
            X_train = X_train.to(device)
            y_train = y_train.to(device)

            y_pred = model(X_train)
            

            # 计算损失
            loss = criterion(y_pred, y_train)  # 预测值与真实值对比
            is_correct = getIsCorrect(y_pred, y_train) # 预测正确返回10，错误返回1
            fitnesses_in_epoch.append(is_correct)

            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            b += 1
        
        fitnesses_each_epoch.append(math.fmean(fitnesses_in_epoch))

    manager.fitnessesResult('TRAIN in batch ' + str(currentBatchNum), fitnesses_each_epoch)

测试输出 #

TEST in type tensor([0]): mean: 10.0
TEST in type tensor([1]): mean: 1.0
TEST in type tensor([2]): mean: 1.0

可见首个标签（tensor([0])）测试表现完美（均值10.0），但后续标签（tensor([1])、tensor([2])）仅达到最低分1.0，完全未习得有效知识。

补充信息 #

• 每个数据集批次约含6000个元素
• 每个数据集训练2个epoch
• 模型层结构：

Conv2d(1, 6, kernel_size=(3, 3), stride=(1, 1))
Conv2d(6, 16, kernel_size=(3, 3), stride=(1, 1))
GRU(400, 32)
RNN(32, 32)
GRU(32, 40)
Dropout(p=0.778, inplace=False)

• 额外配置：

activation functions: ['sigmoid', 'none', 'none', 'sigmoid']
epochs: 2
learn rate: 0.01
optimizer: adam
test size: 0.8

我已尝试调整模型层、优化器、epoch数、学习率、测试集比例等配置，但问题始终存在，仅表现程度略有差异。请问是否是首个数据集批次训练后，模型状态影响了后续新标签的学习？

分析与解决方案 #

你的问题本质是灾难性遗忘，这是增量学习（持续学习）中最常见的问题：模型在学习新任务时，会覆盖之前任务学到的权重，同时单标签批次训练的模式让模型无法建立新标签与旧标签的区分边界。

核心问题点 #

1. 损失函数的局限性：CrossEntropyLoss在单标签批次训练时，模型只会优化当前标签对应的输出维度，完全忽略其他维度的权重维护，导致旧标签的识别能力被破坏，同时无法学会区分新标签与未见过的类别。
2. 优化器重复初始化：每次调用trainModel都重新创建优化器，丢失了之前训练的动量信息，导致权重更新出现突变，无法平滑适配新任务。
3. 单标签训练的固有缺陷：训练新标签时没有旧标签样本作为参考，模型无法建立标签间的区分逻辑，自然无法正确识别新标签。

具体修复方案 #

1. 引入蒸馏损失避免遗忘

训练新标签时，保留模型对旧标签的输出能力，通过蒸馏损失约束新模型的输出与旧模型（训练前的状态）尽可能接近：

import copy
import torch.nn.functional as F

# 在训练循环外保存初始模型状态
old_model = copy.deepcopy(model)
old_model.eval()

# 修改trainModel函数，加入蒸馏损失
def trainModel(params: Parameters, model: CreatedModel, train_loader, currentBatchNum, device, optimizer):
    ce_criterion = nn.CrossEntropyLoss()
    distill_criterion = nn.KLDivLoss()
    fitnesses_each_epoch = []

    for i in range(params.epochs):
        fitnesses_in_epoch = []
        for b, (X_train, y_train) in enumerate(train_loader):
            X_train = X_train.to(device)
            y_train = y_train.to(device)

            y_pred = model(X_train)
            # 分类损失
            ce_loss = ce_criterion(y_pred, y_train)
            # 蒸馏损失：约束旧标签维度的输出与旧模型一致
            with torch.no_grad():
                old_pred = old_model(X_train)
            distill_loss = distill_criterion(F.log_softmax(y_pred[:, :currentBatchNum], dim=1), 
                                            F.softmax(old_pred[:, :currentBatchNum], dim=1))
            # 总损失：平衡分类与蒸馏权重
            loss = ce_loss + 0.3 * distill_loss

            is_correct = getIsCorrect(y_pred, y_train)
            fitnesses_in_epoch.append(is_correct)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        
        fitnesses_each_epoch.append(math.fmean(fitnesses_in_epoch))
    manager.fitnessesResult('TRAIN in batch ' + str(currentBatchNum), fitnesses_each_epoch)

2. 复用优化器，避免重复初始化

将优化器的初始化移到数据集遍历循环外，保证训练过程中优化器的状态持续更新：

# 只初始化一次优化器
optimizer = get_optimizer(params, model=model)

for index, datasetbatch in enumerate(dataset_batches):
    # 传入已初始化的optimizer
    trainModel(params, model, datasetbatch[1], index + 1, device, optimizer)

3. 加入旧标签样本重放（经验回放）

维护一个旧标签样本池，每次训练新标签时混入少量旧标签样本，让模型记住之前的任务：

from torch.utils.data import ConcatDataset
import random

# 初始化样本池，每个标签保存200个样本
sample_pool = {}

for index, datasetbatch in enumerate(dataset_batches):
    current_label = index
    current_dataset = datasetbatch[1].dataset
    # 构建混合数据集：当前标签数据 + 样本池中的旧数据
    mixed_datasets = [current_dataset]
    for label in sample_pool:
        mixed_datasets.append(sample_pool[label])
    mixed_dataset = ConcatDataset(mixed_datasets)
    mixed_loader = DataLoader(mixed_dataset, batch_size=datasetbatch[1].batch_size, shuffle=True)
    
    # 训练混合数据集
    trainModel(params, model, mixed_loader, index + 1, device, optimizer)
    
    # 将当前标签的部分样本加入样本池
    sample_pool[current_label] = random.sample(list(current_dataset), 200)

4. 调整输出层与激活函数

模型最后一层是GRU输出40维，搭配sigmoid激活更适合多标签任务，但你使用的是CrossEntropyLoss（适配单标签分类）。建议：

• 确保输出维度等于总类别数
• 最后一层移除sigmoid激活，CrossEntropyLoss会自动处理Softmax逻辑
• 中间层的sigmoid可替换为ReLU，避免梯度消失问题

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