autogluon.multimodal 1.2.1b20250303__py3-none-any.whl → 1.2.1b20250304__py3-none-any.whl

autogluon/multimodal/optim/losses/softmax_losses.py

@@ -0,0 +1,177 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+try:
+    import torch.distributed.nn
+    from torch import distributed as dist
+
+    has_distributed = True
+except ImportError:
+    has_distributed = False
+
+try:
+    import horovod.torch as hvd
+except ImportError:
+    hvd = None
+
+
+class SoftTargetCrossEntropy(nn.Module):
+    """
+    The soft target CrossEntropy from timm.
+    https://github.com/rwightman/pytorch-image-models/blob/e4360e6125bb0bb4279785810c8eb33b40af3ebd/timm/loss/cross_entropy.py
+    It works under the mixup.
+    It can calculate the crossentropy of input and label with one-hot.
+    """
+
+    def __init__(self):
+        super(SoftTargetCrossEntropy, self).__init__()
+
+    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        loss = torch.sum(-target * F.log_softmax(input, dim=-1), dim=-1)
+        return loss.mean()
+
+
+class MultiNegativesSoftmaxLoss(nn.Module):
+    """
+    This loss expects as input a batch consisting of pairs (a_1, p_1), (a_2, p_2)…, (a_n, p_n) where
+        we assume that (a_i, p_i) are a positive pair and (a_i, p_j) for i!=j a negative pair.
+        For each a_i, it uses all other p_j as negative samples, i.e., for a_i,
+        we have 1 positive example (p_i) and n-1 negative examples (p_j).
+        It then minimizes the negative log-likehood for softmax normalized scores.
+        It can also support gather negatives across processes.
+    """
+
+    def __init__(
+        self,
+        local_loss=False,
+        gather_with_grad=False,
+        cache_labels=False,
+        use_horovod=False,
+    ):
+        """
+        Parameters
+        ----------
+        local_loss
+            Whether to compute the loss only for the current process's samples.
+        gather_with_grad
+            Whether to gather all features with gradients enabled.
+        cache_labels
+            Whether to cache labels for loss in next iterations.
+        use_horovod
+            Whether to use horovod.
+        """
+        super().__init__()
+        self.local_loss = local_loss
+        self.gather_with_grad = gather_with_grad
+        self.cache_labels = cache_labels
+        self.use_horovod = use_horovod
+
+        # cache state
+        self.prev_num_logits = 0
+        self.labels = {}
+
+    def forward(self, features_a, features_b, logit_scale, rank=0, world_size=1):
+        device = features_a.device
+        if world_size > 1:
+            all_features_a, all_features_b = self.gather_features(
+                features_a, features_b, self.local_loss, self.gather_with_grad, rank, world_size, self.use_horovod
+            )
+
+            if self.local_loss:
+                logits_per_a = logit_scale * features_a @ all_features_b.T
+                logits_per_b = logit_scale * features_b @ all_features_a.T
+            else:
+                logits_per_a = logit_scale * all_features_a @ all_features_b.T
+                logits_per_b = logits_per_a.T
+        else:
+            logits_per_a = logit_scale * features_a @ features_b.T
+            logits_per_b = logit_scale * features_b @ features_a.T
+
+        # calculated ground-truth and cache if enabled
+        num_logits = logits_per_a.shape[0]
+        if self.prev_num_logits != num_logits or device not in self.labels:
+            labels = torch.arange(num_logits, device=device, dtype=torch.long)
+            if world_size > 1 and self.local_loss:
+                labels = labels + num_logits * rank
+            if self.cache_labels:
+                self.labels[device] = labels
+                self.prev_num_logits = num_logits
+        else:
+            labels = self.labels[device]
+
+        total_loss = (F.cross_entropy(logits_per_a, labels) + F.cross_entropy(logits_per_b, labels)) / 2
+        return total_loss
+
+    @staticmethod
+    def gather_features(
+        image_features,
+        text_features,
+        local_loss=False,
+        gather_with_grad=False,
+        rank=0,
+        world_size=1,
+        use_horovod=False,
+    ):
+        """
+        Gather features across GPUs.
+
+        Parameters
+        ----------
+        image_features
+            image features of the current process.
+        text_features
+            text features of the current process.
+        local_loss
+            If False, make sure the features on the current GPU have gradients.
+        gather_with_grad
+            Whether to gather all features with gradients enabled.
+        rank
+            Rank of the current process (it should be a number between 0 and world_size-1).
+        world_size
+            Number of processes participating in the job.
+        use_horovod
+            Whether to use horovod.
+
+        Returns
+        -------
+        Gathered image and text features from all processes.
+        """
+        assert (
+            has_distributed
+        ), "torch.distributed did not import correctly, please use a PyTorch version with support."
+        if use_horovod:
+            assert hvd is not None, "Please install horovod"
+            if gather_with_grad:
+                all_image_features = hvd.allgather(image_features)
+                all_text_features = hvd.allgather(text_features)
+            else:
+                with torch.no_grad():
+                    all_image_features = hvd.allgather(image_features)
+                    all_text_features = hvd.allgather(text_features)
+                if not local_loss:
+                    # ensure grads for local rank when all_* features don't have a gradient
+                    gathered_image_features = list(all_image_features.chunk(world_size, dim=0))
+                    gathered_text_features = list(all_text_features.chunk(world_size, dim=0))
+                    gathered_image_features[rank] = image_features
+                    gathered_text_features[rank] = text_features
+                    all_image_features = torch.cat(gathered_image_features, dim=0)
+                    all_text_features = torch.cat(gathered_text_features, dim=0)
+        else:
+            # We gather tensors from all gpus
+            if gather_with_grad:
+                all_image_features = torch.cat(torch.distributed.nn.all_gather(image_features), dim=0)
+                all_text_features = torch.cat(torch.distributed.nn.all_gather(text_features), dim=0)
+            else:
+                gathered_image_features = [torch.zeros_like(image_features) for _ in range(world_size)]
+                gathered_text_features = [torch.zeros_like(text_features) for _ in range(world_size)]
+                dist.all_gather(gathered_image_features, image_features)
+                dist.all_gather(gathered_text_features, text_features)
+                if not local_loss:
+                    # ensure grads for local rank when all_* features don't have a gradient
+                    gathered_image_features[rank] = image_features
+                    gathered_text_features[rank] = text_features
+                all_image_features = torch.cat(gathered_image_features, dim=0)
+                all_text_features = torch.cat(gathered_text_features, dim=0)
+
+        return all_image_features, all_text_features

autogluon/multimodal/optim/losses/structure_loss.py

@@ -0,0 +1,26 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class StructureLoss(nn.Module):
+    """
+    Structure Loss based on https://github.com/DengPingFan/PraNet/blob/master/MyTrain.py
+    The loss represent the weighted IoU loss and binary cross entropy (BCE) loss for the global restriction and local (pixel-level) restriction.
+
+    References:
+        [1] https://arxiv.org/abs/2006.11392
+    """
+
+    def forward(self, input: torch.Tensor, target: torch.Tensor):
+        if input.dim() == 3:
+            input = input.unsqueeze(1)
+        weit = 1 + 5 * torch.abs(F.avg_pool2d(target, kernel_size=31, stride=1, padding=15) - target)
+        wbce = F.binary_cross_entropy_with_logits(input, target, reduce="none")
+        wbce = (weit * wbce).sum(dim=(2, 3)) / weit.sum(dim=(2, 3))
+
+        input = torch.sigmoid(input)
+        inter = ((input * target) * weit).sum(dim=(2, 3))
+        union = ((input + target) * weit).sum(dim=(2, 3))
+        wiou = 1 - (inter + 1) / (union - inter + 1)
+        return (wbce + wiou).mean()

autogluon/multimodal/optim/losses/utils.py

@@ -0,0 +1,313 @@
+import logging
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+from omegaconf import DictConfig, OmegaConf
+from pytorch_metric_learning import distances, losses, miners
+from torch import nn
+from transformers.models.mask2former.modeling_mask2former import Mask2FormerConfig, Mask2FormerLoss
+
+from ...constants import (
+    BINARY,
+    CONTRASTIVE_LOSS,
+    COSINE_SIMILARITY,
+    FEW_SHOT_CLASSIFICATION,
+    MULTI_NEGATIVES_SOFTMAX_LOSS,
+    MULTICLASS,
+    NER,
+    OBJECT_DETECTION,
+    PAIR_MARGIN_MINER,
+    REGRESSION,
+    SEMANTIC_SEGMENTATION,
+)
+from .bce_loss import BBCEWithLogitLoss
+from .focal_loss import FocalLoss
+from .lemda_loss import LemdaLoss
+from .softmax_losses import MultiNegativesSoftmaxLoss, SoftTargetCrossEntropy
+from .structure_loss import StructureLoss
+
+logger = logging.getLogger(__name__)
+
+
+def get_loss_func(
+    problem_type: str,
+    mixup_active: Optional[bool] = None,
+    loss_func_name: Optional[str] = None,
+    config: Optional[DictConfig] = None,
+    **kwargs,
+):
+    """
+    Choose a suitable Pytorch loss module based on the provided problem type.
+
+    Parameters
+    ----------
+    problem_type
+        Type of problem.
+    mixup_active
+        The activation determining whether to use mixup.
+    loss_func_name
+        The name of the function the user wants to use.
+    config
+        The optimization configs containing values such as i.e. optim.loss_func
+        An example purpose of this config here is to pass through the parameters for focal loss, i.e.:
+            alpha = optim.focal_loss.alpha
+    Returns
+    -------
+    A Pytorch loss module.
+    """
+    if problem_type in [BINARY, MULTICLASS]:
+        if mixup_active:
+            loss_func = SoftTargetCrossEntropy()
+        else:
+            if loss_func_name is not None and loss_func_name.lower() == "focal_loss":
+                loss_func = FocalLoss(
+                    alpha=config.focal_loss.alpha,
+                    gamma=config.focal_loss.gamma,
+                    reduction=config.focal_loss.reduction,
+                )
+            else:
+                loss_func = nn.CrossEntropyLoss(label_smoothing=config.label_smoothing)
+                logger.debug(f"loss_func.label_smoothing: {loss_func.label_smoothing}")
+    elif problem_type == REGRESSION:
+        if loss_func_name is not None:
+            if "bcewithlogitsloss" in loss_func_name.lower():
+                loss_func = nn.BCEWithLogitsLoss()
+            else:
+                loss_func = nn.MSELoss()
+        else:
+            loss_func = nn.MSELoss()
+    elif problem_type == NER:
+        loss_func = nn.CrossEntropyLoss(ignore_index=0)
+    elif problem_type in [None, OBJECT_DETECTION, FEW_SHOT_CLASSIFICATION]:
+        return None
+    elif problem_type == SEMANTIC_SEGMENTATION:
+        if "structure_loss" in loss_func_name.lower():
+            loss_func = StructureLoss()
+        elif "balanced_bce" in loss_func_name.lower():
+            loss_func = BBCEWithLogitLoss()
+        elif "mask2former_loss" in loss_func_name.lower():
+            weight_dict = {
+                "loss_cross_entropy": config.mask2former_loss.loss_cross_entropy_weight,
+                "loss_mask": config.mask2former_loss.loss_mask_weight,
+                "loss_dice": config.mask2former_loss.loss_dice_weight,
+            }
+            loss_func = Mask2FormerLoss(
+                config=Mask2FormerConfig(num_labels=kwargs["num_classes"]), weight_dict=weight_dict
+            )
+        else:
+            loss_func = nn.BCEWithLogitsLoss()
+    else:
+        raise NotImplementedError
+
+    return loss_func
+
+
+def get_metric_learning_distance_func(
+    name: str,
+):
+    """
+    Return one pytorch metric learning's distance function based on its name.
+
+    Parameters
+    ----------
+    name
+        distance function name
+
+    Returns
+    -------
+    A distance function from the pytorch metric learning package.
+    """
+    if name.lower() == COSINE_SIMILARITY:
+        return distances.CosineSimilarity()
+    else:
+        raise ValueError(f"Unknown distance measure: {name}")
+
+
+def infer_matcher_loss(data_format: str, problem_type: str):
+    """
+    Infer the loss type to train the matcher.
+
+    Parameters
+    ----------
+    data_format
+        The training data format, e.g., pair or triplet.
+    problem_type
+        Type of problem.
+
+    Returns
+    -------
+    The loss name.
+    """
+    if data_format == "pair":
+        if problem_type is None:
+            return [MULTI_NEGATIVES_SOFTMAX_LOSS]
+        elif problem_type == BINARY:
+            return [CONTRASTIVE_LOSS]
+        elif problem_type == REGRESSION:
+            return ["cosine_similarity_loss"]
+        else:
+            raise ValueError(f"Unsupported data format {data_format} with problem type {problem_type}")
+    elif data_format == "triplet":
+        if problem_type is None:
+            return [MULTI_NEGATIVES_SOFTMAX_LOSS]
+        else:
+            raise ValueError(f"Unsupported data format {data_format} with problem type {problem_type}")
+    else:
+        raise ValueError(f"Unsupported data format: {data_format}")
+
+
+def get_matcher_loss_func(
+    data_format: str,
+    problem_type: str,
+    loss_type: Optional[str] = None,
+    pos_margin: Optional[float] = None,
+    neg_margin: Optional[float] = None,
+    distance_type: Optional[str] = None,
+):
+    """
+    Return a list of pytorch metric learning's loss functions based on their names.
+
+    Parameters
+    ----------
+    data_format
+        The training data format, e.g., pair or triplet.
+    problem_type
+        Type of problem.
+    loss_type
+        The provided loss type.
+    pos_margin
+        The positive margin in computing the metric learning loss.
+    neg_margin
+        The negative margin in computing the metric learning loss.
+    distance_type
+        The distance function type.
+
+    Returns
+    -------
+    A loss function of metric learning.
+    """
+
+    allowable_loss_types = infer_matcher_loss(data_format=data_format, problem_type=problem_type)
+    if loss_type is not None:
+        assert loss_type in allowable_loss_types, f"data format {data_format} can't use loss {loss_type}."
+    else:
+        loss_type = allowable_loss_types[0]
+
+    if loss_type.lower() == CONTRASTIVE_LOSS:
+        return losses.ContrastiveLoss(
+            pos_margin=pos_margin,
+            neg_margin=neg_margin,
+            distance=get_metric_learning_distance_func(distance_type),
+        )
+    elif loss_type.lower() == MULTI_NEGATIVES_SOFTMAX_LOSS:
+        return MultiNegativesSoftmaxLoss(
+            local_loss=True,
+            gather_with_grad=True,
+            cache_labels=False,
+        )
+    else:
+        raise ValueError(f"Unknown metric learning loss: {loss_type}")
+
+
+def get_matcher_miner_func(
+    miner_type: str,
+    pos_margin: float,
+    neg_margin: float,
+    distance_type: str,
+):
+    """
+    Return a pytorch metric learning's miner functions based on their names.
+    The miners are used to mine the positive and negative examples.
+
+    Parameters
+    ----------
+    miner_type
+        The miner function type.
+    pos_margin
+        The positive margin used by the miner function.
+    neg_margin
+        The negative margin used by the miner function.
+    distance_type
+        The distance function type.
+
+    Returns
+    -------
+    A miner function to mine positive and negative samples.
+    """
+    if miner_type.lower() == PAIR_MARGIN_MINER:
+        return miners.PairMarginMiner(
+            pos_margin=pos_margin,
+            neg_margin=neg_margin,
+            distance=get_metric_learning_distance_func(distance_type),
+        )
+    else:
+        raise ValueError(f"Unknown metric learning miner: {miner_type}")
+
+
+def generate_metric_learning_labels(
+    num_samples: int,
+    match_label: int,
+    labels: torch.Tensor,
+):
+    """
+    Generate labels to compute the metric learning loss of one mini-batch.
+    For n samples, it generates 2*n labels since each match has two sides, each of which
+    has one label. If we know the matching label, then it determines the two sides' labels
+    according to whether their label is the matching label. If the matching label is None,
+    it assigns a unique label for each side.
+
+    Parameters
+    ----------
+    num_samples
+        number of samples.
+    match_label
+        The matching label, which can be None.
+    labels
+        The sample labels used in the supervised setting. It's required only when match_label is not None.
+
+    Returns
+    -------
+    The labels used in computing the metric learning loss.
+    """
+    device = labels.device
+    labels_1 = torch.arange(num_samples, device=device)
+
+    if match_label is not None:
+        labels_2 = torch.arange(num_samples, num_samples * 2, device=device)
+        # users need to specify the match_label based on the raw label's semantic meaning.
+        mask = labels == match_label
+        labels_2[mask] = labels_1[mask]
+    else:
+        labels_2 = torch.arange(num_samples, device=device)
+
+    metric_learning_labels = torch.cat([labels_1, labels_2], dim=0)
+
+    return metric_learning_labels
+
+
+def get_aug_loss_func(config: Optional[DictConfig] = None, problem_type: Optional[str] = None):
+    """
+    Return the loss function for lemda augmentation
+
+    Parameters
+    ----------
+    config
+        The optimization configuration.
+    problem_type
+        Problem type (binary, multimclass, or regression)
+
+    Returns
+    -------
+    Augmentation loss function.
+    """
+    loss_func = None
+    if config.lemda.turn_on:
+        loss_func = LemdaLoss(
+            mse_weight=config.lemda.mse_weight,
+            kld_weight=config.lemda.kld_weight,
+            consist_weight=config.lemda.consist_weight,
+            consist_threshold=config.lemda.consist_threshold,
+            problem_type=problem_type,
+        )
+
+    return loss_func

autogluon/multimodal/optim/lr/__init__.py

@@ -0,0 +1 @@
+from .utils import apply_layerwise_lr_decay, apply_single_lr, apply_two_stages_lr, get_lr_scheduler