import torch import torch.nn as nn import torch.nn.functional as F import numpy as np from torch.utils.data import DataLoader from collections import deque from typing import List, Dict, Any, Optional, Union from tqdm import tqdm from dataclasses import dataclass @dataclass class ModalityConfig: name: str modality_id: int class UnifiedMultiModalPreprocessor(nn.Module): def __init__(self, model_dim: int = 2048): super().__init__() self.modality_configs = { 'text': ModalityConfig('text', 0), 'image': ModalityConfig('image', 1), 'audio': ModalityConfig('audio', 2), 'video': ModalityConfig('video', 3) } def process_batch(self, batch_data: Union[torch.Tensor, List[Any]], modality_type: str) -> List[Dict]: processed_segments = [] if modality_type not in self.modality_configs: return processed_segments config = self.modality_configs[modality_type] if isinstance(batch_data, list): # 过滤 None valid_data = [x for x in batch_data if x is not None] if not valid_data: return [] # 假设 list 中全是 Tensor，且维度一致，进行堆叠 # 如果是 list of tensor (B, C, H, W) -> stack -> (B, C, H, W) try: data_tensor = torch.stack(valid_data) except Exception as e: print(f"Error stacking modality data: {e}") return [] elif isinstance(batch_data, torch.Tensor): data_tensor = batch_data else: return [] processed_segments.append({ 'type': modality_type, 'data': data_tensor, 'modality_id': config.modality_id }) return processed_segments class ExperienceReplayBuffer: def __init__(self, max_size: int = 10000): self.buffer = deque(maxlen=max_size) def add(self, sample: Dict[str, Any]): safe_sample = {} for k, v in sample.items(): if isinstance(v, torch.Tensor): safe_sample[k] = v.detach().cpu() elif isinstance(v, list): # 递归处理 list 中的 tensor safe_sample[k] = [x.detach().cpu() if isinstance(x, torch.Tensor) else x for x in v] else: safe_sample[k] = v self.buffer.append(safe_sample) def sample(self, batch_size: int) -> List[Any]: """从buffer中采样""" if not self.buffer: return [] indices = np.random.choice( len(self.buffer), min(len(self.buffer), batch_size), replace=False ) return [self.buffer[i] for i in indices] def __len__(self): return len(self.buffer) def clear(self): """清空buffer""" self.buffer.clear() class EWC: """弹性权重固化 (Elastic Weight Consolidation)""" def __init__( self, model: nn.Module, dataloader: DataLoader, preprocessor: UnifiedMultiModalPreprocessor, importance: float = 1000.0 ): self.model = model self.preprocessor = preprocessor self.importance = importance self.device = next(model.parameters()).device # 冻结当前参数作为参考 self.params = { n: p.clone().detach() for n, p in model.named_parameters() if p.requires_grad } self.fisher = self._compute_fisher(dataloader) def _compute_fisher(self, dataloader: DataLoader) -> Dict[str, torch.Tensor]: """计算Fisher信息矩阵 (使用 Empirical Fisher)""" fisher = { n: torch.zeros_like(p) for n, p in self.model.named_parameters() if p.requires_grad } self.model.eval() num_samples = 0 # 使用 tqdm 稍微简化输出 pbar = tqdm(dataloader, desc="Computing Fisher Matrix", leave=False) for batch in pbar: if batch is None: continue self.model.zero_grad() # 1. 准备文本输入 instruction_ids = batch['instruction'].to(self.device) response_ids = batch['response'].to(self.device) # 拼接: [Instruction, Response] input_ids = torch.cat([instruction_ids, response_ids], dim=1) # 2. 准备多模态输入结构 input_data = {'segments': []} # 处理额外的模态数据 raw_modality_data = batch.get('modality_data') if raw_modality_data is not None: modality_type = batch.get('modality_type', 'image') if isinstance(modality_type, list): modality_type = modality_type[0] mod_segments = self.preprocessor.process_batch(raw_modality_data, modality_type) for seg in mod_segments: seg['data'] = seg['data'].to(self.device) input_data['segments'].append(seg) input_data['segments'].append({ 'type': 'text', 'data': input_ids, 'modality_id': 0 }) output = self.model(input_data) logits = output['logits'] # 4. 计算 Loss (Standard Causal LM Loss) # Shift logits and labels # input_ids: [I1, I2, R1, R2] # labels: [I2, R1, R2, EOS] shift_logits = logits[:, :-1, :].contiguous() shift_labels = input_ids[:, 1:].contiguous() # 创建 Mask: 只在 Response 部分计算梯度 # Instruction 长度 inst_len = instruction_ids.shape[1] loss_mask = torch.ones_like(shift_labels, dtype=torch.float) if inst_len > 1: loss_mask[:, :inst_len-1] = 0.0 # 计算逐个 Token 的 Loss loss_fct = nn.CrossEntropyLoss(reduction='none') loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) # 应用 Mask 并求平均 loss = (loss * loss_mask.view(-1)).sum() / (loss_mask.sum() + 1e-6) # 5. 反向传播累积梯度平方 loss.backward() for n, p in self.model.named_parameters(): if p.grad is not None and n in fisher: fisher[n] += p.grad.detach() ** 2 num_samples += input_ids.size(0) # 平均化 if num_samples > 0: for n in fisher: fisher[n] /= num_samples self.model.train() return fisher def penalty(self, model: Optional[nn.Module] = None) -> torch.Tensor: target_model = model if model is not None else self.model loss = torch.tensor(0.0, device=self.device) for n, p in target_model.named_parameters(): if n in self.params and p.requires_grad: if n in self.fisher: loss += (self.fisher[n] * (p - self.params[n]) ** 2).sum() return self.importance * loss class OnlineEWC(EWC): def __init__( self, model: nn.Module, preprocessor: UnifiedMultiModalPreprocessor, importance: float = 1000.0, gamma: float = 0.9 ): self.model = model self.preprocessor = preprocessor self.importance = importance self.gamma = gamma self.device = next(model.parameters()).device self.params = {} self.fisher = {} self.task_count = 0 def update_fisher(self, dataloader: DataLoader): """更新Fisher信息矩阵""" print(f"Updating Online EWC Fisher Matrix (Task {self.task_count + 1})...") new_fisher = self._compute_fisher(dataloader) if self.task_count == 0: self.fisher = new_fisher else: for n in self.fisher: if n in new_fisher: # 移动平均更新 Fisher 信息 self.fisher[n] = self.gamma * self.fisher[n] + new_fisher[n] # 更新参考参数为当前任务训练后的参数 self.params = { n: p.clone().detach() for n, p in self.model.named_parameters() if p.requires_grad } self.task_count += 1 print(f"Online EWC regularizer updated.") def penalty(self, model: Optional[nn.Module] = None) -> torch.Tensor: """计算EWC惩罚项""" if self.task_count == 0: return torch.tensor(0.0, device=self.device) return super().penalty(model)