| 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) |
|
|
| |
| 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: |
| 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() |
|
|
| |
| 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) |
| latents = rearrange(latents, 'm n l d -> n l (m d)') |
|
|
| |
| |
|
|
| 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): |
| |
| 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)) |
|
|
