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 # noqa: F401 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}