| import logging |
|
|
| import numpy as np |
| import torch |
| import torch.distributed as dist |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from entmax import entmax_bisect |
|
|
| from .dfr import DFRLoss |
|
|
| logger = logging.getLogger("glint") |
|
|
|
|
| def pad_and_gather(tensor): |
| local_size = torch.tensor(tensor.size(), device=tensor.device) |
|
|
| all_sizes = [torch.zeros_like(local_size) for _ in range(dist.get_world_size())] |
| dist.all_gather(all_sizes, local_size) |
|
|
| max_size = torch.stack(all_sizes).max(dim=0)[0] |
|
|
| padded_tensor = torch.zeros(max_size.tolist(), device=tensor.device) |
| padded_tensor[: local_size[0]] = tensor |
|
|
| gathered_tensors = dist.nn.all_gather(padded_tensor) |
|
|
| gathered_tensors = [g[: s[0]] for g, s in zip(gathered_tensors, all_sizes)] |
|
|
| gathered_tensors = torch.cat(gathered_tensors, dim=0) |
|
|
| return gathered_tensors |
|
|
|
|
| def pad_and_gather_2nd_dim(tensor, mask): |
| local_size = torch.tensor(tensor.size(), device=tensor.device) |
|
|
| if local_size[1] == 0: |
| local_size = torch.tensor([local_size[0], 1, local_size[2]], device=tensor.device) |
| tensor = torch.zeros( |
| (local_size[0], 1, local_size[2]), |
| device=tensor.device, |
| requires_grad=True, |
| ) |
| mask = torch.zeros((local_size[0], 1), device=mask.device, dtype=torch.bool) |
|
|
| all_sizes = [torch.zeros_like(local_size) for _ in range(dist.get_world_size())] |
| dist.all_gather(all_sizes, local_size) |
|
|
| max_len = max([i[1].item() for i in all_sizes]) |
|
|
| padded_tensor = torch.zeros((local_size[0], max_len, local_size[2]), device=tensor.device) |
| padded_tensor[:, : local_size[1]] = tensor |
|
|
| padded_mask = torch.zeros((local_size[0], max_len), device=mask.device) |
| padded_mask[:, : local_size[1]] = mask |
|
|
| gathered_tensors = dist.nn.all_gather(padded_tensor) |
| gathered_masks = dist.nn.all_gather(padded_mask) |
|
|
| gathered_tensors = torch.cat(gathered_tensors, dim=0) |
| gathered_masks = torch.cat(gathered_masks, dim=0) |
|
|
| return gathered_tensors, gathered_masks |
|
|
|
|
| def _build_projection(embed_dim, hidden_dim, mode="mlp"): |
| if mode == "identity": |
| return nn.Identity() |
| if mode == "linear": |
| return nn.Linear(embed_dim, hidden_dim) |
| if mode == "mlp": |
| return nn.Sequential( |
| nn.Linear(embed_dim, hidden_dim), |
| nn.GELU(), |
| nn.Linear(hidden_dim, embed_dim), |
| ) |
| raise ValueError(f"Unknown projection mode: {mode}") |
|
|
|
|
| class GatedConfidenceShift(nn.Module): |
| def __init__( |
| self, |
| embed_dim=768, |
| gate_hidden_dim=768, |
| sem_hidden_dim=768, |
| gate_proj_mode="mlp", |
| sem_proj_mode="identity", |
| gate_log_temperature=2.302585, |
| shift_w_init=0.0, |
| shift_w_no_weight_decay=True, |
| ): |
| super().__init__() |
|
|
| self.text_gate_proj = _build_projection(embed_dim, gate_hidden_dim, gate_proj_mode) |
| self.vision_gate_proj = _build_projection(embed_dim, gate_hidden_dim, gate_proj_mode) |
| self.gate_log_temperature = nn.Parameter(torch.FloatTensor([gate_log_temperature])) |
|
|
| self.text_sem_proj = _build_projection(embed_dim, sem_hidden_dim, sem_proj_mode) |
| self.vision_sem_proj = _build_projection(embed_dim, sem_hidden_dim, sem_proj_mode) |
|
|
| self.shift_w = nn.Parameter(torch.tensor(float(shift_w_init))) |
| if shift_w_no_weight_decay: |
| self.shift_w._no_weight_decay = True |
|
|
| def _compute_gate_evidence(self, gate): |
| w = entmax_bisect(gate, alpha=2.0, n_iter=50, dim=-1) |
| return (w * gate).sum(dim=-1) |
|
|
| def forward(self, text_features, vision_tokens, need_attn_weights=False, repeat=True, **kwargs): |
| if repeat: |
| query = text_features.unsqueeze(0).expand(vision_tokens.shape[0], text_features.shape[0], -1) |
| else: |
| assert text_features.dim() == 3 |
| query = text_features |
|
|
| text_gate = F.normalize(self.text_gate_proj(query), p=2, dim=-1) |
| vision_gate = F.normalize(self.vision_gate_proj(vision_tokens), p=2, dim=-1) |
| gate_scores = torch.bmm(text_gate, vision_gate.permute(0, 2, 1)) |
|
|
| gate = torch.sigmoid(gate_scores * self.gate_log_temperature.exp()) |
| weight = gate |
|
|
| gate_evidence = self._compute_gate_evidence(gate) |
| shift = self.shift_w * gate_evidence |
|
|
| sem_text_norm = F.normalize(self.text_sem_proj(query), p=2, dim=-1) |
| sem_patches_norm = F.normalize(self.vision_sem_proj(vision_tokens), p=2, dim=-1) |
|
|
| aggregated = torch.bmm(weight, sem_patches_norm) |
| aggregated_norm = F.normalize(aggregated, p=2, dim=-1) |
| cos_sim = (sem_text_norm * aggregated_norm).sum(dim=-1) |
| logits = (cos_sim + shift).T |
|
|
| if need_attn_weights: |
| sem_scores = torch.bmm(sem_text_norm, sem_patches_norm.permute(0, 2, 1)) |
| if "loss_temperature" in kwargs: |
| scaled_sem = (sem_scores / kwargs["loss_temperature"]).sigmoid() |
| elif "logit_scale" in kwargs: |
| scaled_sem = (sem_scores * kwargs["logit_scale"] + kwargs.get("logit_bias", 0.0)).sigmoid() |
| else: |
| scaled_sem = sem_scores.sigmoid() |
|
|
| similarity_scores = [scaled_sem * weight] |
| _comp = {"gate": gate, "semantic": scaled_sem, "weight": weight} |
| else: |
| similarity_scores = None |
| _comp = None |
|
|
| return logits, similarity_scores, _comp |
|
|
|
|
| def multi_positive_nce_loss( |
| logits, |
| group_map, |
| temperature, |
| additional_neg_mask=None, |
| ): |
| scaled_logits = torch.exp(logits / temperature) |
| np_total = group_map.size(0) |
| pos_logits = scaled_logits[torch.arange(np_total), group_map] |
|
|
| row_loss = get_row_loss( |
| scaled_logits[:np_total], |
| pos_logits, |
| ) |
|
|
| neg_mask = torch.ones_like(scaled_logits) |
|
|
| if additional_neg_mask is not None: |
| neg_mask[-len(additional_neg_mask) :, :] = additional_neg_mask |
|
|
| neg_mask[torch.arange(np_total), group_map] = 0 |
|
|
| column_loss = get_col_loss( |
| scaled_logits, |
| pos_logits, |
| neg_mask, |
| group_map, |
| ) |
|
|
| losses = {} |
| losses["per_text_loss"] = row_loss.mean() |
| losses["per_image_loss"] = column_loss.mean() |
|
|
| losses["mpnce_loss"] = (losses["per_text_loss"] + losses["per_image_loss"]) / 2 |
|
|
| losses["loss"] = losses["mpnce_loss"] |
|
|
| return losses |
|
|
|
|
| def multi_positive_nce_local_loss( |
| logits_per_text, |
| logits_per_image, |
| group_map, |
| local_group_map, |
| local_group_map_per_image, |
| temperature, |
| local_rank, |
| B_local, |
| additional_neg_mask=None, |
| ): |
| logits_per_text = torch.exp(logits_per_text / temperature) |
| pos_logits_per_text = logits_per_text[torch.arange(logits_per_text.size(0)), local_group_map] |
|
|
| row_loss = get_row_loss( |
| logits_per_text, |
| pos_logits_per_text, |
| ) |
|
|
| logits_per_image = torch.exp(logits_per_image / temperature) |
|
|
| local_start_idx = torch.where(group_map == local_rank * B_local)[0][0] |
| local_end_idx = local_start_idx + len(local_group_map_per_image) |
|
|
| pos_text_indices = torch.LongTensor(range(local_start_idx, local_end_idx)) |
|
|
| pos_logits_per_image = logits_per_image[pos_text_indices, local_group_map_per_image] |
|
|
| neg_mask = torch.ones_like(logits_per_image) |
|
|
| if additional_neg_mask is not None: |
| neg_mask[-len(additional_neg_mask) :, :] = additional_neg_mask |
|
|
| neg_mask[pos_text_indices, local_group_map_per_image] = 0 |
|
|
| column_loss = get_col_loss( |
| logits_per_image, |
| pos_logits_per_image, |
| neg_mask, |
| local_group_map_per_image, |
| ) |
|
|
| losses = {} |
| losses["per_text_loss"] = row_loss.mean() |
| losses["per_image_loss"] = column_loss.mean() |
|
|
| losses["mpnce_loss"] = (losses["per_text_loss"] + losses["per_image_loss"]) / 2 |
|
|
| losses["loss"] = losses["mpnce_loss"] |
| return losses |
|
|
|
|
| def get_row_loss( |
| logits: torch.Tensor, |
| pos_logits: torch.Tensor, |
| eps: float = 1e-8, |
| ): |
| row_sum_logits = logits.sum(dim=1) |
| p_row = pos_logits / (row_sum_logits + eps) |
|
|
| return -torch.log(p_row + eps) |
|
|
|
|
| def get_col_loss( |
| logits: torch.Tensor, |
| pos_logits: torch.Tensor, |
| neg_mask: torch.Tensor, |
| group_map: torch.Tensor, |
| eps: float = 1e-8, |
| ): |
| neg_logits = logits * neg_mask |
| sum_neg_logits = neg_logits.sum(dim=0) |
| sum_neg_logits = sum_neg_logits[group_map] |
| p_column = pos_logits / (pos_logits + sum_neg_logits + eps) |
|
|
| return -torch.log(p_column + eps) |
|
|
|
|
| class LossFunctions(nn.Module): |
| def __init__(self, loss_cfg, **kwargs): |
| super().__init__() |
| self.loss_ratio = dict() |
| self.loss_dict = nn.ModuleDict() |
| for loss_type, ratio in zip(loss_cfg["apply"], loss_cfg["ratio"]): |
| logger.info(f"Build {loss_type} loss function ...") |
| if loss_cfg[loss_type] is None: |
| loss_cfg[loss_type] = dict() |
| if torch.distributed.is_available() and torch.distributed.is_initialized(): |
| loss_cfg[loss_type]["rank"] = torch.distributed.get_rank() |
| loss_cfg[loss_type]["world_size"] = torch.distributed.get_world_size() |
| self.loss_dict[loss_type] = LOSS_REGISTRY[loss_type](**loss_cfg[loss_type]) |
| self.loss_ratio[loss_type] = ratio |
|
|
|
|
| class AlignmentLoss(nn.Module): |
| def __init__( |
| self, |
| vl_logit_type="gate_confidence_shift", |
| vl_logit_kwargs={}, |
| loss_type="mpnce", |
| loss_kwargs={}, |
| local_loss=False, |
| **kwargs, |
| ): |
| super().__init__() |
|
|
| self.vl_logit_type = vl_logit_type |
| self.init_vl_logit_kwargs(vl_logit_kwargs) |
|
|
| self.loss_type = loss_type |
| self.init_loss_kwargs(loss_kwargs) |
|
|
| self.local_loss = local_loss |
|
|
| def forward( |
| self, |
| enc_sents, |
| vision_tokens, |
| forward_text_model=None, |
| ddp_gather=True, |
| need_attn_weights=False, |
| compute_loss=True, |
| text_feats_list=None, |
| neg_enc_sents=None, |
| neg_text_feats_list=None, |
| **kwargs, |
| ): |
| outputs = {} |
|
|
| B_local = len(vision_tokens) |
| local_rank = dist.get_rank() if (ddp_gather and dist.is_initialized()) else 0 |
|
|
| text_features, group_map, local_text_features, local_group_map = self.get_text_features( |
| enc_sents, forward_text_model, ddp_gather, text_feats_list |
| ) |
|
|
| local_vision_tokens = vision_tokens |
| if ddp_gather and dist.is_initialized(): |
| vision_tokens = torch.cat(dist.nn.all_gather(vision_tokens), dim=0) |
|
|
| if not compute_loss: |
| logits_outputs = self.compute_logits(text_features, vision_tokens, need_attn_weights) |
|
|
| logits = logits_outputs["logits"] |
|
|
| if self.loss_type == "mpnce": |
| logits = logits / self.loss_temperature.exp() |
|
|
| logits_outputs["logits"] = logits |
| outputs.update(logits_outputs) |
| return outputs |
|
|
| if neg_enc_sents or neg_text_feats_list: |
| hidden_dim = vision_tokens.shape[-1] |
|
|
| ( |
| neg_text_features, |
| neg_text_mask, |
| local_neg_text_features, |
| local_neg_text_mask, |
| ) = self.get_neg_text_features( |
| neg_enc_sents, |
| forward_text_model, |
| hidden_dim, |
| ddp_gather, |
| neg_text_feats_list, |
| device=vision_tokens.device, |
| ) |
| else: |
| ( |
| neg_text_features, |
| neg_text_mask, |
| local_neg_text_features, |
| local_neg_text_mask, |
| ) = (None, None, None, None) |
|
|
| if not (self.local_loss and self.training): |
| logits_outputs = self.compute_logits(text_features, vision_tokens, need_attn_weights) |
| logits = logits_outputs["logits"] |
|
|
| if neg_text_features is not None: |
| neg_logits_outputs = self.compute_logits( |
| neg_text_features, vision_tokens, need_attn_weights, repeat=False |
| ) |
|
|
| logits = torch.cat([logits, neg_logits_outputs["logits"]], dim=0) |
|
|
| outputs["logits"] = logits |
|
|
| losses = self.compute_loss( |
| logits, |
| group_map, |
| additional_neg_mask=neg_text_mask, |
| text_features=text_features, |
| ) |
|
|
| else: |
| logits_outputs_per_text = self.compute_logits(local_text_features, vision_tokens, need_attn_weights) |
| logits_per_text = logits_outputs_per_text["logits"] |
| outputs["logits_per_text"] = logits_per_text |
|
|
| logits_outputs_per_image = self.compute_logits(text_features, local_vision_tokens, need_attn_weights) |
| logits_per_image = logits_outputs_per_image["logits"] |
|
|
| if local_neg_text_features is not None: |
| neg_logits_outputs = self.compute_logits( |
| local_neg_text_features, |
| local_vision_tokens, |
| need_attn_weights, |
| repeat=False, |
| ) |
|
|
| logits_per_image = torch.cat([logits_per_image, neg_logits_outputs["logits"]], dim=0) |
|
|
| outputs["logits_per_image"] = logits_per_image |
| local_group_map_per_image = local_group_map - local_rank * B_local |
|
|
| losses = self.compute_local_loss( |
| logits_per_text, |
| logits_per_image, |
| group_map, |
| local_group_map, |
| local_group_map_per_image, |
| local_rank, |
| B_local, |
| additional_neg_mask=local_neg_text_mask, |
| text_features=text_features, |
| ) |
|
|
| outputs["losses"] = losses |
| return outputs |
|
|
| def get_feats_list(self, text_feats_list, ddp_gather=True): |
| text_features = list() |
| group_list = list() |
|
|
| B_local = sum([len(i) for i in text_feats_list]) |
| local_rank = dist.get_rank() if (ddp_gather and dist.is_initialized()) else 0 |
|
|
| local_index = 0 |
| for i, features in enumerate(text_feats_list): |
| for j, feats in enumerate(features): |
| text_features.append(feats) |
| group_list.extend([local_index + local_rank * B_local] * feats.size(0)) |
| local_index += 1 |
|
|
| return text_features, group_list |
|
|
| def get_text_features( |
| self, |
| enc_sents, |
| forward_text_model, |
| ddp_gather=True, |
| text_feats_list=None, |
| ): |
| if text_feats_list is not None: |
| text_features, group_list = self.get_feats_list(text_feats_list, ddp_gather) |
| else: |
| text_features = list() |
| group_list = list() |
|
|
| B_local = len(enc_sents) |
| local_rank = dist.get_rank() if (ddp_gather and dist.is_initialized()) else 0 |
|
|
| for i, kp in enumerate(enc_sents): |
| feats = forward_text_model(kp) |
|
|
| feat = feats["text_features"] |
| text_features.append(feat) |
|
|
| global_index = i + local_rank * B_local |
| group_list.extend([global_index] * feat.size(0)) |
|
|
| text_features = torch.cat(text_features, dim=0) |
| group_map = torch.tensor(group_list, device=text_features.device) |
|
|
| local_text_features = text_features |
| local_group_map = group_map |
|
|
| if ddp_gather and dist.is_initialized(): |
| text_features = pad_and_gather(text_features) |
|
|
| group_map = pad_and_gather(group_map) |
| group_map = group_map.long() |
|
|
| return text_features, group_map, local_text_features, local_group_map |
|
|
| def get_neg_feats_list(self, neg_text_feats_list, hidden_dim=None, device=None): |
| text_features = list() |
| for i, features in enumerate(neg_text_feats_list): |
| for j, feats in enumerate(features): |
| text_features.append(feats) |
|
|
| return text_features |
|
|
| def get_neg_text_features( |
| self, |
| neg_enc_sents, |
| forward_text_model, |
| hidden_dim, |
| ddp_gather=True, |
| neg_text_feats_list=None, |
| device=None, |
| ): |
| if neg_text_feats_list is not None: |
| neg_feats_list = self.get_neg_feats_list(neg_text_feats_list, hidden_dim, device) |
| else: |
| neg_feats_list = list() |
|
|
| for i, kp in enumerate(neg_enc_sents): |
| if isinstance(kp["input_ids"], torch.Tensor): |
| feats = forward_text_model(kp) |
| feat = feats["text_features"] |
| else: |
| feat = torch.zeros((0, hidden_dim)).to(device) |
|
|
| neg_feats_list.append(feat) |
|
|
| neg_text_features = torch.nn.utils.rnn.pad_sequence(neg_feats_list, batch_first=True, padding_value=0) |
|
|
| neg_text_mask = torch.zeros( |
| len(neg_feats_list), |
| neg_text_features.size(1), |
| dtype=torch.bool, |
| device=neg_text_features.device, |
| ) |
| for i, feat in enumerate(neg_feats_list): |
| neg_text_mask[i, : feat.size(0)] = True |
|
|
| local_neg_text_features = neg_text_features |
| local_neg_text_mask = neg_text_mask |
|
|
| if ddp_gather and dist.is_initialized(): |
| neg_text_features, neg_text_mask = pad_and_gather_2nd_dim(neg_text_features, neg_text_mask) |
|
|
| local_neg_text_mask = local_neg_text_mask.T |
| neg_text_mask = neg_text_mask.T |
|
|
| return ( |
| neg_text_features, |
| neg_text_mask, |
| local_neg_text_features, |
| local_neg_text_mask, |
| ) |
|
|
| def init_vl_logit_kwargs(self, vl_logit_kwargs): |
| self.vl_logit_kwargs = {} |
| if self.vl_logit_type != "gate_confidence_shift": |
| raise NotImplementedError(f"Unsupported vl_logit_type: {self.vl_logit_type}") |
| vl_logit_kwargs.pop("attn_temperature", None) |
| self.sga = GatedConfidenceShift(**vl_logit_kwargs) |
|
|
| def compute_logits(self, text_features, vision_tokens, need_attn_weights, repeat=True): |
| outputs = {} |
|
|
| sim_scale_kwargs = {} |
| if hasattr(self, "loss_temperature"): |
| sim_scale_kwargs["loss_temperature"] = self.loss_temperature.exp() |
|
|
| logits, attn_weights_list, sim_components = self.sga( |
| text_features, |
| vision_tokens, |
| need_attn_weights, |
| repeat=repeat, |
| **sim_scale_kwargs, |
| **self.vl_logit_kwargs, |
| ) |
| outputs["logits"] = logits |
| outputs["attn_weights"] = attn_weights_list |
| if sim_components is not None: |
| outputs["sim_components"] = sim_components |
| return outputs |
|
|
| def init_loss_kwargs(self, loss_kwargs): |
| self.loss_kwargs = {} |
| if self.loss_type != "mpnce": |
| raise NotImplementedError(f"Unsupported loss_type: {self.loss_type}") |
| self.loss_temperature = nn.Parameter(torch.FloatTensor([np.log(loss_kwargs["temperature"])])) |
|
|
| def compute_loss(self, logits, group_map, additional_neg_mask=None, **kwargs): |
| return multi_positive_nce_loss( |
| logits, |
| group_map, |
| additional_neg_mask=additional_neg_mask, |
| temperature=self.loss_temperature.exp(), |
| **self.loss_kwargs, |
| ) |
|
|
| def compute_local_loss( |
| self, |
| logits_per_text, |
| logits_per_image, |
| group_map, |
| local_group_map, |
| local_group_map_per_image, |
| local_rank, |
| B_local, |
| additional_neg_mask=None, |
| **kwargs, |
| ): |
| return multi_positive_nce_local_loss( |
| logits_per_text, |
| logits_per_image, |
| group_map, |
| local_group_map, |
| local_group_map_per_image, |
| temperature=self.loss_temperature.exp(), |
| local_rank=local_rank, |
| B_local=B_local, |
| additional_neg_mask=additional_neg_mask, |
| **self.loss_kwargs, |
| ) |
|
|
|
|
| LOSS_REGISTRY = {"AlignmentLoss": AlignmentLoss, "DFRLoss": DFRLoss} |
|
|
