import copy from collections import OrderedDict import torch import torch.nn as nn import torch.nn.functional as F from einops import rearrange class PretrainingWrapper(nn.Module): """Wrapper that owns DINO/SimCLR pretraining routines and delegates feature extraction to a backbone model. Note: KoLeo (`koleo_loss`) and DGCL (`dgcl_loss`) remain on the backbone. """ def __init__(self, backbone, simclr=False, simclr_temp=0.1, dino=False, dino_student_temp=0.1, dino_teacher_temp=0.04, dino_center_momentum=0.9, dino_momentum=0.996): super().__init__() self.backbone = backbone self.simclr = simclr self.simclr_temp = simclr_temp self.dino = dino self.dino_student_temp = dino_student_temp self.dino_teacher_temp = dino_teacher_temp self.dino_center_momentum = dino_center_momentum self.dino_momentum = dino_momentum self.use_micro = getattr(backbone, 'use_micro', False) self.bidi = getattr(backbone, 'bidi', False) self.cycle_len = getattr(backbone, 'cycle_len', 180) self.use_transformer = getattr(backbone, 'use_transformer', False) # teacher modules (set by init_teacher) self.teacher_encoder = None self.teacher_fuser = None self.teacher_proj = None self.teacher_sequence_encoder = None self.teacher_contr_head = None if self.dino: # try to infer seq_dim from backbone.proj try: seq_dim = self.backbone.proj.out_features except Exception: seq_dim = 512 self.register_buffer('dino_center_f', torch.zeros(1, seq_dim)) self.register_buffer('dino_center_r', torch.zeros(1, seq_dim)) def forward(self, sig, mask=None, demo=None): if self.dino: return self._forward_dino(sig, mask) if self.simclr: return self._forward_simclr(sig, mask) return self.backbone(sig, mask=mask, demo=demo) def __getattr__(self, name): """Delegate unknown attributes to backbone to preserve legacy access patterns like model.encoder/model.fuser/model.revin. """ try: return super().__getattr__(name) except AttributeError as exc: if self.backbone is not None and hasattr(self.backbone, name): return getattr(self.backbone, name) raise exc def state_dict(self, *args, **kwargs): """Expose original backbone key names (without `backbone.` prefix).""" sd = super().state_dict(*args, **kwargs) out = OrderedDict() for k, v in sd.items(): if k.startswith('backbone.'): out[k[len('backbone.'):]] = v else: out[k] = v return out def load_state_dict(self, state_dict, strict=True): """Accept both legacy keys and wrapped keys.""" own_keys = super().state_dict().keys() converted = OrderedDict() for k, v in state_dict.items(): if k in own_keys: converted[k] = v elif f'backbone.{k}' in own_keys: converted[f'backbone.{k}'] = v else: converted[k] = v return super().load_state_dict(converted, strict=strict) def init_teacher(self): if not self.dino: return self.teacher_encoder = copy.deepcopy(self.backbone.encoder) self.teacher_fuser = copy.deepcopy(self.backbone.fuser) self.teacher_proj = copy.deepcopy(self.backbone.proj) if hasattr(self.backbone, 'proj') else None self.teacher_sequence_encoder = copy.deepcopy(self.backbone.sequence_encoder) if hasattr(self.backbone, 'sequence_encoder') else None self.teacher_contr_head = copy.deepcopy(self.backbone.contr_head) if hasattr(self.backbone, 'contr_head') else None self._freeze_teacher() def _freeze_teacher(self): for module in [self.teacher_encoder, self.teacher_fuser, self.teacher_proj, self.teacher_sequence_encoder, self.teacher_contr_head]: if module is None: continue for p in module.parameters(): p.requires_grad = False module.eval() # helper that mirrors previous behaviour in-class def _forward_features_with_modules(self, inputs, mask, encoder, fuser, proj, sequence_encoder): modalities_to_use = sorted(list(inputs.keys())) m = modalities_to_use[0] N, L = mask.shape if mask is not None: nonmask_idx = torch.nonzero(mask.view(-1,)).squeeze(-1) latents_shared, latents_specific, _ = encoder(inputs, modalities_to_use) latents_fused, _ = fuser(latents_shared, latents_specific) latents = torch.stack(latents_fused, dim=0) # (M, N, L, D) latents = rearrange(latents, 'm n l d -> n l (m d)') # if getattr(self.backbone, 'use_micro', False): # return None, latents_micro latents = proj(latents) latents = latents.mean(1) dummy = torch.zeros((N * L, latents.shape[-1]), device=latents.device, dtype=latents.dtype) dummy.scatter_(dim=0, index=nonmask_idx.unsqueeze(-1).repeat(1, latents.shape[-1]), src=latents) latents = dummy latents = latents.view(N, -1, latents.shape[-1]) latents = sequence_encoder(latents, mask=mask) return latents, None def _get_cycle_values(self, xf, mask, mode='last'): if hasattr(self.backbone, '_get_cycle_values'): return self.backbone._get_cycle_values(xf, mask, mode=mode) if mask is not None: xf = xf.masked_fill(~mask.unsqueeze(-1), 0.0) min_length = mask.sum(dim=1).min() num_seg = min_length // self.cycle_len if num_seg < 1: seq1_last = xf[torch.arange(xf.shape[0], device=xf.device), mask.sum(dim=1) - 1] return seq1_last.unsqueeze(1) seq1 = xf.view(xf.shape[0], -1, self.cycle_len, xf.shape[-1]).mean(2) return seq1 def _simclr_representation(self, sig, mask): feats = self.backbone.forward_features(sig, mask=mask) x_macro, _ = feats if self.bidi: xf, xr = x_macro xf = self._get_cycle_values(xf, mask) xr = self._get_cycle_values(xr, mask) return xf, xr else: return self._get_cycle_values(x_macro, mask) def _simclr_representation_simple(self, sig, mask): # alias kept for clarity return self._simclr_representation(sig, mask) def _simclr_loss(self, z1, z2): z1 = F.normalize(z1, p=2, dim=-1) z2 = F.normalize(z2, p=2, dim=-1) z = torch.cat([z1, z2], dim=0) logits = torch.matmul(z, z.t()) / (self.simclr_temp + 1e-8) logits = logits - torch.eye(logits.shape[0], device=logits.device, dtype=logits.dtype) * 1e9 batch_size = z1.shape[0] targets = torch.arange(2 * batch_size, device=logits.device) targets = (targets + batch_size) % (2 * batch_size) return F.cross_entropy(logits, targets) def _forward_simclr(self, sig, mask): if mask is None: raise ValueError("SimCLR pretraining requires a valid mask.") if isinstance(sig, (tuple, list)): if len(sig) != 2: raise ValueError("SimCLR expects exactly two views in a list/tuple.") sig1, sig2 = sig elif isinstance(sig, dict) and 'view1' in sig and 'view2' in sig: sig1, sig2 = sig['view1'], sig['view2'] else: raise ValueError("SimCLR expects inputs as (view1, view2) or a dict with 'view1'/'view2'.") if self.bidi: z1_f, z1_r = self._simclr_representation(sig1, mask) z2_f, z2_r = self._simclr_representation(sig2, mask) z1_f = self.backbone.contr_head(z1_f) z1_r = self.backbone.contr_head(z1_r) z2_f = self.backbone.contr_head(z2_f) z2_r = self.backbone.contr_head(z2_r) num_seg = z1_f.shape[1] loss_fwd, loss_bwd = 0.0, 0.0 for s in range(num_seg): loss_fwd += self._simclr_loss(z1_f[:,s], z2_f[:,s]) loss_bwd += self._simclr_loss(z1_r[:,s], z2_r[:,s]) loss = (loss_fwd + loss_bwd) / 2 else: z1 = self._simclr_representation(sig1, mask) z2 = self._simclr_representation(sig2, mask) loss = self._simclr_loss(self.backbone.contr_head(z1), self.backbone.contr_head(z2)) return loss def _dino_representation(self, sig, mask, use_teacher=False): if use_teacher: feats = self._forward_features_with_modules(sig, mask, self.teacher_encoder, self.teacher_fuser, self.teacher_proj, self.teacher_sequence_encoder) else: feats = self.backbone.forward_features(sig, mask=mask) x_macro, _ = feats if self.bidi: xf, xr = x_macro xf = self._get_cycle_values(xf, mask) xr = self._get_cycle_values(xr, mask) return xf, xr else: return self._get_cycle_values(x_macro, mask) def _forward_dino(self, sig, mask): if self.simclr: raise ValueError("DINO and SimCLR cannot be enabled at the same time.") if self.teacher_encoder is None: raise ValueError("DINO teacher modules are not initialized. Call init_teacher() after loading checkpoints.") if mask is None: raise ValueError("DINO pretraining requires a valid mask.") if isinstance(sig, (tuple, list)): if len(sig) != 2: raise ValueError("DINO expects exactly two views in a list/tuple.") sig1, sig2 = sig elif isinstance(sig, dict) and 'view1' in sig and 'view2' in sig: sig1, sig2 = sig['view1'], sig['view2'] else: raise ValueError("DINO expects inputs as (view1, view2) or a dict with 'view1'/'view2'.") z1_f_s, z1_r_s = self._dino_representation(sig1, mask, use_teacher=False) z2_f_s, z2_r_s = self._dino_representation(sig2, mask, use_teacher=False) z1_f_s = self.backbone.contr_head(z1_f_s) z1_r_s = self.backbone.contr_head(z1_r_s) z2_f_s = self.backbone.contr_head(z2_f_s) z2_r_s = self.backbone.contr_head(z2_r_s) with torch.no_grad(): z1_f_t, z1_r_t = self._dino_representation(sig1, mask, use_teacher=True) z2_f_t, z2_r_t = self._dino_representation(sig2, mask, use_teacher=True) z1_f_t = self.teacher_contr_head(z1_f_t) z1_r_t = self.teacher_contr_head(z1_r_t) z2_f_t = self.teacher_contr_head(z2_f_t) z2_r_t = self.teacher_contr_head(z2_r_t) self._update_dino_center(torch.cat([z1_f_t.mean(1), z2_f_t.mean(1)], dim=0), torch.cat([z1_r_t.mean(1), z2_r_t.mean(1)], dim=0)) num_seg = z2_f_s.shape[1] loss_fwd, loss_bwd = 0.0, 0.0 for s in range(num_seg): loss_fwd += self.dino_loss(z1_f_s[:,s], z2_f_t[:,s], fwd=True) loss_fwd += self.dino_loss(z2_f_s[:,s], z1_f_t[:,s], fwd=True) loss_bwd += self.dino_loss(z1_r_s[:,s], z2_r_t[:,s], fwd=False) loss_bwd += self.dino_loss(z2_r_s[:,s], z1_r_t[:,s], fwd=False) loss = (loss_fwd + loss_bwd) / 2 / 2 return loss def dino_loss(self, student_out, teacher_out, fwd=False): student_out = student_out / (self.dino_student_temp + 1e-8) center = self.dino_center_f if fwd else self.dino_center_r teacher_out = F.softmax((teacher_out - center) / (self.dino_teacher_temp + 1e-8), dim=-1) return torch.sum(-teacher_out * F.log_softmax(student_out, dim=-1), dim=-1).mean() @torch.no_grad() def _update_dino_center(self, teacher_out_f, teacher_out_r): batch_center_f = teacher_out_f.mean(dim=0, keepdim=True) batch_center_r = teacher_out_r.mean(dim=0, keepdim=True) self.dino_center_f.mul_(self.dino_center_momentum).add_(batch_center_f * (1.0 - self.dino_center_momentum)) self.dino_center_r.mul_(self.dino_center_momentum).add_(batch_center_r * (1.0 - self.dino_center_momentum)) @torch.no_grad() def update_teacher(self, momentum=None): if not self.dino or self.teacher_encoder is None: return m = self.dino_momentum if momentum is None else momentum pairs = [ (self.backbone.encoder, self.teacher_encoder), (self.backbone.fuser, self.teacher_fuser), ] if self.teacher_proj is not None: pairs.append((self.backbone.proj, self.teacher_proj)) if self.teacher_sequence_encoder is not None: pairs.append((self.backbone.sequence_encoder, self.teacher_sequence_encoder)) pairs.append((self.backbone.contr_head, self.teacher_contr_head)) for student, teacher in pairs: for param_s, param_t in zip(student.parameters(), teacher.parameters()): param_t.data.mul_(m).add_(param_s.data * (1.0 - m))