| 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) |
|
|
