import logging import torch import torch.nn as nn import torch.nn.functional as F from .vision_encoders import build_vision_encoder logger = logging.getLogger("glint") def build_projector(config): if config.model_type == "mlp": model = MLP(config) else: raise NotImplementedError() return model def get_activation(activation_name): activation_map = {"silu": nn.SiLU(), "relu": nn.ReLU(), "gelu": nn.GELU(), "tanh": nn.Tanh()} activation_name = activation_name.lower() if activation_name not in activation_map: raise ValueError(f"Unknown activation: {activation_name}. Available: {list(activation_map.keys())}") return activation_map[activation_name] class MLP(nn.Module): def __init__(self, config): super().__init__() self.config = config activation = get_activation(getattr(config, "activation", "silu")) dropout = getattr(config, "dropout", 0.0) num_layers = getattr(config, "num_layers", 3) layers = [] layers.append(nn.Linear(config.input_dim, config.hidden_dim)) layers.append(activation) if dropout > 0: layers.append(nn.Dropout(dropout)) for _ in range(num_layers - 2): layers.append(nn.Linear(config.hidden_dim, config.hidden_dim)) layers.append(activation) if dropout > 0: layers.append(nn.Dropout(dropout)) layers.append(nn.Linear(config.hidden_dim, config.out_dim)) self.mlp = nn.Sequential(*layers) def forward(self, x): return self.mlp(x) class DFRLoss(nn.Module): def __init__( self, use_token_wise=True, **kwargs, ): super().__init__() self.use_token_wise = use_token_wise def forward( self, proj_student_features, teacher_features, ): outputs = dict() if isinstance(proj_student_features, list): total_loss = 0.0 n = 0 for ps, ts in zip(proj_student_features, teacher_features): if isinstance(ps, dict): for k in ps.keys(): total_loss += self._compute_loss(ps[k], ts[k]) n += 1 else: total_loss += self._compute_loss(ps, ts) n += 1 loss = total_loss / max(n, 1) elif isinstance(proj_student_features, dict): total_loss = 0 num_pairs = len(proj_student_features) for key in proj_student_features.keys(): proj_student = proj_student_features[key] teacher = teacher_features[key] total_loss += self._compute_loss(proj_student, teacher) loss = total_loss / num_pairs else: loss = self._compute_loss(proj_student_features, teacher_features) losses = { "dfr_loss": loss, "loss": loss, } outputs["losses"] = losses return outputs def _compute_loss(self, proj_student, teacher): proj_student = F.normalize(proj_student, p=2, dim=-1) teacher = F.normalize(teacher, p=2, dim=-1) if self.use_token_wise: similarity = torch.sum(proj_student * teacher, dim=-1) loss = (1.0 - similarity).mean() else: proj_student_global = proj_student.mean(dim=1) teacher_global = teacher.mean(dim=1) similarity = torch.sum(proj_student_global * teacher_global, dim=-1) loss = (1.0 - similarity).mean() loss = loss.mean() return loss class DFRModule(nn.Module): def __init__(self, dfr_config): super().__init__() self.config = dfr_config self.use_cache = getattr(dfr_config, "use_cache", False) if not self.use_cache: self.vision_teacher_model = build_vision_encoder(dfr_config.vfm_teacher_config) else: logger.info("DFRModule: use_cache=True, skipping teacher model loading") self.layer_mapping = dfr_config.layer_mapping student_layers = self.layer_mapping["student_layers"] teacher_layers = self.layer_mapping["teacher_layers"] self._validate_layer_mapping(student_layers, teacher_layers) self.num_layer_pairs = len(student_layers) self._build_projectors(dfr_config.projector_config, student_layers, teacher_layers) def _validate_layer_mapping(self, student_layers, teacher_layers): if len(student_layers) != len(teacher_layers): raise ValueError( f"student_layer and teacher_layer must have same length. " f"Got {len(student_layers)} and {len(teacher_layers)}" ) def _build_projectors(self, projector_config, student_layers, teacher_layers): use_multiple_projectors = self.layer_mapping.get("use_multiple_projectors", False) if use_multiple_projectors and self.num_layer_pairs > 1: self.projector = nn.ModuleDict() for s_layer, t_layer in zip(student_layers, teacher_layers): projector_key = f"projector_{s_layer}_{t_layer}" self.projector[projector_key] = build_projector(projector_config) self.use_multiple_projectors = True else: self.projector = build_projector(projector_config) self.use_multiple_projectors = False def forward(self, pixel_values, student_hidden_states, cached_teacher_features=None): student_layers = self.layer_mapping["student_layers"] teacher_layers = self.layer_mapping["teacher_layers"] if cached_teacher_features is not None: teacher_hidden_states = cached_teacher_features else: with torch.no_grad(): teacher_outputs = self.vision_teacher_model(pixel_values, output_hidden_states=True) results = { "teacher_features": {}, "proj_student_features": {}, } for s_layer, t_layer in zip(student_layers, teacher_layers): student_feat = student_hidden_states[s_layer] if cached_teacher_features is not None: teacher_feat = teacher_hidden_states[t_layer] teacher_feat = teacher_feat.to(dtype=student_feat.dtype) else: teacher_feat = teacher_outputs.hidden_states[t_layer] proj_student_feat = self._project_student_features(student_feat, s_layer, t_layer) results["teacher_features"][f"layer_{s_layer}_{t_layer}"] = teacher_feat results["proj_student_features"][f"layer_{s_layer}_{t_layer}"] = proj_student_feat return results def _project_student_features(self, student_feat, s_layer, t_layer): if self.use_multiple_projectors: projector_key = f"projector_{s_layer}_{t_layer}" return self.projector[projector_key](student_feat) else: return self.projector(student_feat)