| import logging |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from .align_layers import build_align_layer |
| from .common_layers import BasePreTrainedModel |
| from .configuration_glint import GLINTConfig |
| from .losses import LossFunctions |
| from .text_encoders import build_text_encoder |
| from .vision_encoders import build_vision_encoder |
|
|
| logger = logging.getLogger("glint") |
|
|
|
|
| class GLINTModel(BasePreTrainedModel): |
| config_class = GLINTConfig |
| |
| _keys_to_ignore_on_load_unexpected = [ |
| r"^text_model\.", |
| r"^align_layer\.", |
| r"^loss_fns\.", |
| r"^dfr_module\.", |
| r"^vision_layer_norm\.", |
| r"^text_layer_norm\.", |
| ] |
|
|
| def build_vision_model(self, config): |
| vision_config = config.vision_config |
| vision_model = build_vision_encoder(vision_config) |
|
|
| return vision_model |
|
|
| def build_text_model(self, config): |
| text_config = config.text_config |
| text_model = build_text_encoder(text_config) |
|
|
| return text_model |
|
|
| def build_align_layer_model(self, config): |
| align_layer_config = config.align_layer_config |
| align_layer = build_align_layer(align_layer_config) |
|
|
| return align_layer |
|
|
| def build_dfr_module(self, config): |
| from .dfr import DFRModule |
|
|
| dfr_module = DFRModule(config.dfr_config) |
| return dfr_module |
|
|
| def __init__(self, config: GLINTConfig): |
| super().__init__(config) |
|
|
| logger.info("Build vision model ...") |
| self.vision_model = self.build_vision_model(config) |
|
|
| if getattr(config, "vision_only", False): |
| logger.info("vision_only=True: skipping text model, align layer, losses and DFR module.") |
| return |
|
|
| logger.info("Build text model ...") |
| self.text_model = self.build_text_model(config) |
|
|
| logger.info("Build align layer model ...") |
| self.align_layer = self.build_align_layer_model(config) |
|
|
| logger.info("Build loss functions ...") |
| loss_cfg = config.kwargs["loss"] |
|
|
| self.loss_fns = LossFunctions(loss_cfg) |
|
|
| if hasattr(config, "dfr_config") and config.dfr_config is not None: |
| logger.info("Build dfr module ...") |
| self.dfr_module = self.build_dfr_module(config) |
|
|
| self.vision_layer_norm = ( |
| nn.LayerNorm(config.align_layer_config.hidden_size) if config.use_vision_layer_norm else None |
| ) |
| self.text_layer_norm = nn.LayerNorm(config.text_config.hidden_size) if config.use_text_layer_norm else None |
|
|
| self.compute_logits_type = config.kwargs.get("compute_logits_type") |
| self.module_to_update = config.kwargs.get("module_to_update") |
|
|
| def forward_vision_model(self, pixel_values, output_hidden_states=False): |
| vision_model_type = self.config.vision_config.model_type |
|
|
| vision_hidden_states = None |
|
|
| use_cls_token = getattr(self.config.vision_config, "use_cls_token", False) |
|
|
| if vision_model_type == "dinov2": |
| vision_outputs = self.vision_model(pixel_values, output_hidden_states=output_hidden_states) |
| vision_tokens = vision_outputs["last_hidden_state"] |
| if not use_cls_token: |
| vision_tokens = vision_tokens[:, 1:, :] |
| vision_hidden_states = vision_outputs.get("hidden_states", None) |
|
|
| elif vision_model_type == "dinov3_vit": |
| vision_outputs = self.vision_model(pixel_values, output_hidden_states=output_hidden_states) |
| original_output = vision_outputs["last_hidden_state"] |
| patch_tokens = original_output[:, 1 + self.config.vision_config.num_register_tokens :, :] |
| if use_cls_token: |
| cls_token = original_output[:, :1, :] |
| vision_tokens = torch.cat([cls_token, patch_tokens], dim=1) |
| else: |
| vision_tokens = patch_tokens |
|
|
| vision_hidden_states = vision_outputs.get("hidden_states", None) |
|
|
| elif vision_model_type == "vjepa2": |
| vision_outputs = self.vision_model( |
| pixel_values, skip_predictor=True, output_hidden_states=output_hidden_states |
| ) |
| vision_tokens = vision_outputs["last_hidden_state"] |
| vision_hidden_states = vision_outputs.get("hidden_states", None) |
| elif vision_model_type == "vjepa2.1": |
| vision_outputs = self.vision_model(pixel_values, output_hidden_states=output_hidden_states) |
| vision_tokens = vision_outputs["last_hidden_state"] |
| vision_hidden_states = vision_outputs.get("hidden_states", None) |
| else: |
| raise NotImplementedError |
|
|
| outputs = {} |
|
|
| vision_tokens = self.align_layer(vision_tokens) |
|
|
| if self.vision_layer_norm is not None: |
| vision_tokens = self.vision_layer_norm(vision_tokens) |
|
|
| outputs["vision_tokens"] = vision_tokens |
| outputs["vision_hidden_states"] = vision_hidden_states |
|
|
| return outputs |
|
|
| def forward_text_model(self, encoded_input): |
| text_outputs = {} |
|
|
| text_model_type = self.config.text_config.model_type |
| if text_model_type == "mpnet": |
| model_output = self.text_model(encoded_input["input_ids"], attention_mask=encoded_input["attention_mask"]) |
| token_embeddings = model_output[0] |
|
|
| else: |
| raise NotImplementedError |
|
|
| if self.config.text_config.use_cls_token: |
| text_features = token_embeddings[:, 0, :] |
| else: |
| input_mask_expanded = ( |
| encoded_input["attention_mask"].unsqueeze(-1).expand(token_embeddings.size()).to(token_embeddings.dtype) |
| ) |
| text_features = torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp( |
| input_mask_expanded.sum(1), min=1e-9 |
| ) |
|
|
| if self.text_layer_norm is not None: |
| text_features = self.text_layer_norm(text_features) |
|
|
| text_outputs["text_features"] = text_features |
| text_outputs["text_features_norm"] = F.normalize(text_features, p=2, dim=1) |
|
|
| return text_outputs |
|
|
| def forward( |
| self, |
| pixel_values=None, |
| enc_sents=None, |
| neg_enc_sents=None, |
| return_loss=True, |
| forward_accum_outputs=False, |
| forward_accum_loss=False, |
| **kwargs, |
| ): |
| if forward_accum_outputs: |
| return self.forward_accum_outputs( |
| pixel_values=pixel_values, |
| enc_sents=enc_sents, |
| neg_enc_sents=neg_enc_sents, |
| **kwargs, |
| ) |
|
|
| if forward_accum_loss: |
| return self.forward_accum_loss(**kwargs) |
|
|
| if getattr(self.config, "vision_only", False): |
| return self.vision_model(pixel_values) |
|
|
| outputs = {} |
| use_dfr_module = hasattr(self, "dfr_module") and self.dfr_module is not None |
| vision_outputs = self.forward_vision_model(pixel_values, output_hidden_states=use_dfr_module) |
|
|
| if use_dfr_module: |
| dfr_outputs = self.dfr_module( |
| pixel_values=pixel_values, |
| student_hidden_states=vision_outputs["vision_hidden_states"], |
| cached_teacher_features=kwargs.get("cached_teacher_features"), |
| ) |
| outputs.update(dfr_outputs) |
|
|
| outputs.update(vision_outputs) |
|
|
| if return_loss: |
| loss = 0 |
| losses = {} |
|
|
| for loss_type, loss_fn in self.loss_fns.loss_dict.items(): |
| loss_ratio = self.loss_fns.loss_ratio[loss_type] |
|
|
| if loss_fn.__class__.__name__ == "AlignmentLoss": |
| loss_outputs = loss_fn( |
| enc_sents=enc_sents, |
| neg_enc_sents=neg_enc_sents, |
| vision_tokens=outputs["vision_tokens"], |
| forward_text_model=self.forward_text_model, |
| ) |
| _losses = loss_outputs["losses"] |
| loop_loss = _losses.pop("loss") |
| losses.update(_losses) |
|
|
| elif loss_fn.__class__.__name__ == "DFRLoss": |
| loss_outputs = loss_fn( |
| proj_student_features=outputs["proj_student_features"], |
| teacher_features=outputs["teacher_features"], |
| ) |
| _losses = loss_outputs["losses"] |
| loop_loss = _losses.pop("loss") |
| losses.update(_losses) |
|
|
| else: |
| raise NotImplementedError |
|
|
| loss += loop_loss * loss_ratio |
|
|
| losses["loss"] = loss |
|
|
| outputs["loss"] = loss |
| outputs["losses"] = losses |
|
|
| return outputs |
|
|
| def compute_logits( |
| self, |
| pixel_values, |
| enc_sents, |
| **kwargs, |
| ): |
| vision_outputs = self.forward_vision_model(pixel_values) |
|
|
| outputs = {} |
|
|
| if self.compute_logits_type == "vl_alignment": |
| loss_fn = self.loss_fns.loss_dict["AlignmentLoss"] |
|
|
| splited_sentences = [ |
| { |
| "input_ids": enc_sents[0]["input_ids"][i : i + 1], |
| "attention_mask": enc_sents[0]["attention_mask"][i : i + 1], |
| } |
| for i in range(enc_sents[0]["input_ids"].size(0)) |
| ] |
|
|
| loss_outputs = loss_fn( |
| splited_sentences, |
| vision_outputs["vision_tokens"], |
| self.forward_text_model, |
| ddp_gather=False, |
| need_attn_weights=True, |
| compute_loss=False, |
| ) |
| outputs.update(loss_outputs) |
|
|
| similarity_scores = loss_outputs["attn_weights"][0] |
| if getattr(self.config.vision_config, "use_cls_token", False): |
| similarity_scores = similarity_scores[:, :, 1:] |
| outputs["similarity_scores"] = similarity_scores |
|
|
| if "sim_components" in loss_outputs: |
| comp = loss_outputs["sim_components"] |
| if getattr(self.config.vision_config, "use_cls_token", False): |
| comp = {k: v[:, :, 1:] for k, v in comp.items()} |
| outputs["sim_components"] = comp |
|
|
| logits = loss_outputs["logits"] |
| logits = logits.T |
|
|
| else: |
| raise NotImplementedError |
|
|
| outputs["logits"] = logits |
| return outputs |
|
|
| def forward_accum_outputs( |
| self, |
| pixel_values, |
| enc_sents=None, |
| neg_enc_sents=None, |
| **kwargs, |
| ): |
| use_dfr_module = hasattr(self, "dfr_module") and self.dfr_module is not None |
|
|
| outputs = {} |
|
|
| vision_outputs = self.forward_vision_model(pixel_values, output_hidden_states=use_dfr_module) |
| outputs.update(vision_outputs) |
|
|
| if use_dfr_module: |
| dfr_outputs = self.dfr_module( |
| pixel_values=pixel_values, |
| student_hidden_states=vision_outputs["vision_hidden_states"], |
| cached_teacher_features=kwargs.get("cached_teacher_features"), |
| ) |
| outputs.update(dfr_outputs) |
|
|
| text_outputs = [] |
| for kp in enc_sents: |
| text_outputs.append(self.forward_text_model(kp)["text_features"]) |
|
|
| text_outputs = {"text_features": text_outputs} |
| outputs.update(text_outputs) |
|
|
| if neg_enc_sents is not None: |
| negative_text_outputs = [] |
| for kp in neg_enc_sents: |
| if isinstance(kp["input_ids"], torch.Tensor): |
| negative_text_outputs.append(self.forward_text_model(kp)["text_features"]) |
| else: |
| negative_text_outputs.append( |
| torch.zeros((0, self.config.text_config.hidden_size)).to(pixel_values.device) |
| ) |
|
|
| negative_text_outputs = {"negative_text_features": negative_text_outputs} |
| outputs.update(negative_text_outputs) |
|
|
| if "ids_mask" in kwargs: |
| outputs["ids_mask"] = kwargs["ids_mask"] |
| outputs["mask"] = kwargs["mask"] |
|
|
| return outputs |
|
|
| def forward_accum_loss( |
| self, |
| vision_tokens, |
| text_features, |
| negative_text_features=None, |
| **kwargs, |
| ): |
| outputs = {} |
| losses = {} |
| loss = 0 |
|
|
| for loss_type, loss_fn in self.loss_fns.loss_dict.items(): |
| loss_ratio = self.loss_fns.loss_ratio[loss_type] |
|
|
| if loss_fn.__class__.__name__ == "AlignmentLoss": |
| loss_outputs = loss_fn( |
| enc_sents=None, |
| neg_enc_sents=None, |
| vision_tokens=vision_tokens, |
| forward_text_model=None, |
| text_feats_list=text_features, |
| neg_text_feats_list=negative_text_features, |
| ) |
| _losses = loss_outputs["losses"] |
| loop_loss = _losses.pop("loss") |
| losses.update(_losses) |
| elif loss_fn.__class__.__name__ == "DFRLoss": |
| loss_outputs = loss_fn( |
| proj_student_features=kwargs.get("proj_student_features"), |
| teacher_features=kwargs.get("teacher_features"), |
| ) |
| _losses = loss_outputs["losses"] |
| loop_loss = _losses.pop("loss") |
| losses.update(_losses) |
| else: |
| raise NotImplementedError |
|
|
| loss += loop_loss * loss_ratio |
|
|
| losses["loss"] = loss |
|
|
| outputs["loss"] = loss |
| outputs["losses"] = losses |
|
|
| return outputs |
|
|
