Blame for losses.py

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}