拒绝“狗熊掰棒子”！用 EWC (Elastic Weight Consolidation) 彻底终结 AI 的灾难性遗忘

/* by 01130.hk - online tools website : 01130.hk/zh/alldns.html */ import torch import torch.nn as nn import torch.optim as optim import matplotlib.pyplot as plt import copy # =========================== # 1. 定义一个简单的神经网络 # =========================== class SimpleNet(nn.Module): def __init__(self): super(SimpleNet, self).__init__() # 为了演示，我们用一个小型的网络 self.fc1 = nn.Linear(10, 20) self.fc2 = nn.Linear(20, 20) self.fc3 = nn.Linear(20, 2) # 输出2类 def forward(self, x): x = torch.relu(self.fc1(x)) x = torch.relu(self.fc2(x)) return self.fc3(x) # =========================== # 2. EWC 核心类 (The Magic) # =========================== class EWC: def __init__(self, model, dataset): self.model = model self.dataset = dataset # 存储旧任务的最优参数 (theta_A*) self.params = {n: p.data.clone() for n, p in self.model.named_parameters()} # 存储每个参数的重要性 (Fisher Information) self.fisher = self._calculate_fisher() def _calculate_fisher(self): fisher = {} # 初始化 Fisher 矩阵为 0 for n, p in self.model.named_parameters(): fisher[n] = torch.zeros_like(p.data) self.model.eval() criterion = nn.CrossEntropyLoss() # 遍历数据计算梯度的平方 # 这里的逻辑是：梯度越大，说明参数稍微一动 Loss 变化就大 -> 参数越重要 for input_data, target in self.dataset: self.model.zero_grad() output = self.model(input_data.unsqueeze(0)) # batch_size=1 loss = criterion(output, target.unsqueeze(0)) loss.backward() for n, p in self.model.named_parameters(): if p.grad is not None: # Fisher 近似等于 梯度的平方 fisher[n] += p.grad.data ** 2 # 归一化 for n in fisher: fisher[n] /= len(self.dataset) return fisher def penalty(self, new_model): loss = 0 for n, p in new_model.named_parameters(): # EWC 公式: Sum( F * (new_theta - old_theta)^2 ) _loss = self.fisher[n] * (p - self.params[n]) ** 2 loss += _loss.sum() return loss # =========================== # 3. 模拟训练流程 # =========================== def get_data(task_id): # 模拟数据：Task 1 输入全1，Task 2 输入全0 if task_id == 1: return [(torch.ones(10), torch.tensor(0)) for _ in range(100)] else: return [(torch.zeros(10), torch.tensor(1)) for _ in range(100)] # 实例化模型 model = SimpleNet() optimizer = optim.SGD(model.parameters(), lr=0.1) criterion = nn.CrossEntropyLoss() print(">>> 开始训练 任务 A (识别全1向量)") data_a = get_data(1) for epoch in range(5): for x, y in data_a: optimizer.zero_grad() loss = criterion(model(x.unsqueeze(0)), y.unsqueeze(0)) loss.backward() optimizer.step() print("任务 A 训练完成。保存 EWC 状态...") # --- 关键步骤：计算 Task A 的重要性权重 --- ewc = EWC(model, data_a) print(">>> 开始训练 任务 B (识别全0向量)，同时开启 EWC 保护") data_b = get_data(2) ewc_lambda = 1000 # 惩罚力度，越大越照顾旧任务 for epoch in range(5): total_loss = 0 for x, y in data_b: optimizer.zero_grad() # 1. 计算新任务的 Loss loss_b = criterion(model(x.unsqueeze(0)), y.unsqueeze(0)) # 2. 计算 EWC 惩罚项 (旧任务的记忆) loss_ewc = ewc.penalty(model) # 3. 总 Loss loss = loss_b + (ewc_lambda * loss_ewc) loss.backward() optimizer.step() total_loss += loss.item() print(f"Epoch {epoch}: Loss = {total_loss:.4f}") # =========================== # 4. 验证结果 # =========================== model.eval() test_a = model(torch.ones(10).unsqueeze(0)) test_b = model(torch.zeros(10).unsqueeze(0)) print("\n=== 最终测试 ===") print(f"Task A (应为类别0): 预测概率 {torch.softmax(test_a, dim=1).detach().numpy()}") print(f"Task B (应为类别1): 预测概率 {torch.softmax(test_b, dim=1).detach().numpy()}")