持续学习突破：DeepSeek灾难性遗忘解决方案

pip install torch torchvision matplotlib numpy

import torch
import torchvision
from torchvision import transforms

# 任务1：MNIST
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
train_dataset_task1 = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset_task1 = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)

# 任务2：Fashion MNIST
train_dataset_task2 = torchvision.datasets.FashionMNIST(root='./data', train=True, download=True, transform=transform)
test_dataset_task2 = torchvision.datasets.FashionMNIST(root='./data', train=False, download=True, transform=transform)

import torch.nn as nn
import torch.nn.functional as F

class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

def compute_fisher(model, data_loader, num_samples=100):
    model.eval()
    fisher = {name: torch.zeros(params.shape).to(params.device) for name, params in model.named_parameters()}
    
    for i, (data, _) in enumerate(data_loader):
        if i >= num_samples:
            break
        data = data.to(device)
        output = model(data)
        loss = output.mean()
        loss.backward()
        
        for name, param in model.named_parameters():
            fisher[name] += param.grad ** 2
    
    for name in fisher:
        fisher[name] /= num_samples
    return fisher

def ewc_loss(model, fisher, previous_params, lambda_ewc=0.1):
    loss = 0.0
    for name, param in model.named_parameters():
        _loss = fisher[name] * (param - previous_params[name]) ** 2
        loss += _loss.sum()
    return lambda_ewc * loss

class ReplayBuffer:
    def __init__(self, capacity=1000):
        self.capacity = capacity
        self.buffer = []
    
    def add(self, data, labels):
        if len(self.buffer) >= self.capacity:
            self.buffer = self.buffer[:self.capacity // 2] + list(zip(data, labels))
        else:
            self.buffer.extend(zip(data, labels))
    
    def sample(self, batch_size):
        indices = torch.randint(len(self.buffer), size=(batch_size,))
        samples = [self.buffer[i] for i in indices]
        data, labels = zip(*samples)
        return torch.stack(data), torch.tensor(labels)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 初始化模型
model = SimpleCNN().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 任务1训练
train_loader_task1 = torch.utils.data.DataLoader(train_dataset_task1, batch_size=64, shuffle=True)
test_loader_task1 = torch.utils.data.DataLoader(test_dataset_task1, batch_size=1000, shuffle=False)

model.train()
for epoch in range(5):
    for batch_idx, (data, target) in enumerate(train_loader_task1):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
    print(f"Task 1 Epoch {epoch+1}, Loss: {loss.item():.4f}")

# 计算Fisher信息矩阵
fisher_task1 = compute_fisher(model, train_loader_task1)
previous_params_task1 = {name: param.clone().detach() for name, param in model.named_parameters()}

# 任务2训练
train_loader_task2 = torch.utils.data.DataLoader(train_dataset_task2, batch_size=64, shuffle=True)
test_loader_task2 = torch.utils.data.DataLoader(test_dataset_task2, batch_size=1000, shuffle=False)

replay_buffer = ReplayBuffer(capacity=1000)
replay_data, replay_labels = next(iter(train_loader_task1))
replay_buffer.add(replay_data, replay_labels)

model.train()
for epoch in range(5):
    for batch_idx, (data, target) in enumerate(train_loader_task2):
        # 添加回放数据
        replay_data, replay_labels = replay_buffer.sample(32)
        combined_data = torch.cat([data, replay_data])
        combined_labels = torch.cat([target, replay_labels])
        
        combined_data, combined_labels = combined_data.to(device), combined_labels.to(device)
        optimizer.zero_grad()
        output = model(combined_data)
        loss = F.nll_loss(output, combined_labels)
        
        # 添加EWC损失
        ewc = ewc_loss(model, fisher_task1, previous_params_task1)
        total_loss = loss + ewc
        
        total_loss.backward()
        optimizer.step()
    print(f"Task 2 Epoch {epoch+1}, Loss: {loss.item():.4f}, EWC Loss: {ewc.item():.4f}")

def evaluate(model, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.nll_loss(output, target, reduction='sum').item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()
    test_loss /= len(test_loader.dataset)
    accuracy = 100. * correct / len(test_loader.dataset)
    return test_loss, accuracy

# 评估任务1
loss_task1, acc_task1 = evaluate(model, test_loader_task1)
print(f"Task 1 Accuracy: {acc_task1:.2f}%")

# 评估任务2
loss_task2, acc_task2 = evaluate(model, test_loader_task2)
print(f"Task 2 Accuracy: {acc_task2:.2f}%")