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

autogluon/multimodal/optim/lr/utils.py

@@ -0,0 +1,332 @@
+import logging
+import re
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+from torch import nn, optim
+
+from ..utils import get_weight_decay_param_names
+from .lr_schedulers import (
+    get_cosine_schedule_with_warmup,
+    get_linear_schedule_with_warmup,
+    get_polynomial_decay_schedule_with_warmup,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def get_lr_scheduler(
+    optimizer: optim.Optimizer,
+    num_max_steps: int,
+    num_warmup_steps: int,
+    lr_schedule: str,
+    end_lr: Union[float, int],
+):
+    """
+    Get the learning rate scheduler from its name. Here we use our defined learning rate
+    scheduler instead of those imported from "transformers" because we want to support
+    Pytorch lightning's "ddp_spawn" training strategy.
+
+    Parameters
+    ----------
+    optimizer
+        A Pytorch optimizer.
+    num_max_steps
+        Number of maximum training steps.
+    num_warmup_steps
+        Number of steps to do learning rate warmup.
+    lr_schedule
+        Name of the learning rate scheduler.
+    end_lr
+        The final learning rate after decay.
+
+    Returns
+    -------
+    A learning rate scheduler.
+    """
+    if lr_schedule == "cosine_decay":
+        scheduler = get_cosine_schedule_with_warmup(
+            optimizer=optimizer,
+            num_warmup_steps=num_warmup_steps,
+            num_training_steps=num_max_steps,
+        )
+    elif lr_schedule == "polynomial_decay":
+        scheduler = get_polynomial_decay_schedule_with_warmup(
+            optimizer=optimizer,
+            num_warmup_steps=num_warmup_steps,
+            num_training_steps=num_max_steps,
+            lr_end=end_lr,
+            power=1,
+        )
+    elif lr_schedule == "linear_decay":
+        scheduler = get_linear_schedule_with_warmup(
+            optimizer=optimizer,
+            num_warmup_steps=num_warmup_steps,
+            num_training_steps=num_max_steps,
+        )
+    elif lr_schedule == "multi_step":
+        # TODO: add milestones, gamma into hyperparameters
+        scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer=optimizer, milestones=[30, 55], gamma=0.1)
+    else:
+        raise ValueError(f"unknown lr schedule: {lr_schedule}")
+
+    return scheduler
+
+
+def apply_single_lr(
+    model: nn.Module,
+    lr: float,
+    weight_decay: float,
+    return_params: Optional[bool] = True,
+    peft: Optional[str] = None,
+    trainable_param_names: Optional[List] = None,
+    exclude_keys: Optional[List] = None,
+):
+    """
+    Set to use a single learning rate for all parameters. Layer normalization parameters and other
+    layers' bias parameters don't use weight decay.
+
+    Parameters
+    ----------
+    model
+        A Pytorch model.
+    lr
+        Learning rate.
+    weight_decay
+        Weight decay.
+    return_params
+        Whether to return parameters or their names. If you want to double-check
+        whether the learning rate setup is as expected, you can set "return_params=False",
+        and print the layer names along with their learning rates through
+        "print("Param groups = %s" % json.dumps(optimizer_grouped_parameters, indent=2))".
+    peft
+        Efficient finetuning strategy. It will only finetune part of the parameters
+    trainable_param_names
+        A list of trainable parameters. (Optional)
+    exclude_keys
+        A list of keys to be excluded.
+
+    Returns
+    -------
+    The grouped parameters or their names.
+    """
+    decay_param_names = get_weight_decay_param_names(model)
+    decay_grad_param_names = []
+    no_decay_grad_param_names = []
+
+    for name, param in model.named_parameters():
+        if exclude_keys and any([exc in name for exc in exclude_keys]):
+            continue
+
+        if (
+            peft is not None
+            and trainable_param_names
+            and not any([re.match(trainable_param_name, name) for trainable_param_name in trainable_param_names])
+        ):
+            param.requires_grad = False
+
+        if not param.requires_grad:
+            continue  # frozen weights
+
+        if name in decay_param_names:
+            if return_params:
+                decay_grad_param_names.append(param)
+            else:
+                decay_grad_param_names.append(name)
+
+        else:
+            if return_params:
+                no_decay_grad_param_names.append(param)
+            else:
+                no_decay_grad_param_names.append(name)
+
+    optimizer_grouped_parameters = [
+        {
+            "params": decay_grad_param_names,
+            "weight_decay": weight_decay,
+            "lr": lr,
+        },
+        {
+            "params": no_decay_grad_param_names,
+            "weight_decay": 0.0,
+            "lr": lr,
+        },
+    ]
+    return optimizer_grouped_parameters
+
+
+def apply_two_stages_lr(
+    model: nn.Module,
+    lr: float,
+    lr_mult: Union[float, int],
+    weight_decay: float,
+    return_params: Optional[bool] = True,
+    exclude_keys: Optional[List] = None,
+):
+    """
+    Set up the pretrained backbone to use a smaller learning rate (lr * lr_mult).
+    The newly added head layers use the normal learning rate (lr).
+    Layer normalization parameters and other layers' bias parameters don't use weight decay.
+
+    Parameters
+    ----------
+    model
+        A Pytorch model.
+    lr
+        The learning rate.
+    lr_mult
+        The multiplier (0, 1) to scale down the learning rate.
+    weight_decay
+        Weight decay.
+    return_params
+        return_params
+        Whether to return parameters or their names. If you want to double-check
+        whether the learning rate setup is as expected, you can set "return_params=False",
+        and print the layer names along with their learning rates through
+        "print("Param groups = %s" % json.dumps(optimizer_grouped_parameters, indent=2))".
+    exclude_keys
+        A list of keys to be excluded.
+
+    Returns
+    -------
+    The grouped parameters or their names.
+    """
+    decay_param_names = get_weight_decay_param_names(model)
+
+    optimizer_grouped_parameters = [
+        {
+            "params": [
+                p if return_params else n
+                for n, p in model.named_parameters()
+                if n in decay_param_names
+                and not any(bb in n for bb in model.head_layer_names)
+                and (not exclude_keys or not any([exc in n for exc in exclude_keys]))
+            ],
+            "weight_decay": weight_decay,
+            "lr": lr,
+        },
+        {
+            "params": [
+                p if return_params else n
+                for n, p in model.named_parameters()
+                if n not in decay_param_names
+                and not any(bb in n for bb in model.head_layer_names)
+                and (not exclude_keys or not any([exc in n for exc in exclude_keys]))
+            ],
+            "weight_decay": 0.0,
+            "lr": lr,
+        },
+        {
+            "params": [
+                p if return_params else n
+                for n, p in model.named_parameters()
+                if n in decay_param_names
+                and any(bb in n for bb in model.head_layer_names)
+                and (not exclude_keys or not any([exc in n for exc in exclude_keys]))
+            ],
+            "weight_decay": weight_decay,
+            "lr": lr * lr_mult,
+        },
+        {
+            "params": [
+                p if return_params else n
+                for n, p in model.named_parameters()
+                if n not in decay_param_names
+                and any(bb in n for bb in model.head_layer_names)
+                and (not exclude_keys or not any([exc in n for exc in exclude_keys]))
+            ],
+            "weight_decay": 0.0,
+            "lr": lr * lr_mult,
+        },
+    ]
+
+    return optimizer_grouped_parameters
+
+
+def apply_layerwise_lr_decay(
+    model: nn.Module,
+    lr: float,
+    lr_decay: float,
+    weight_decay: float,
+    peft: Optional[str] = None,
+    trainable_param_names: Optional[List] = None,
+    exclude_keys: Optional[List] = None,
+):
+    """
+    Assign monotonically decreasing learning rates for layers from the output end to the input end.
+    The intuition behind is that later layers are more task-related compared to the early layers.
+    Layer normalization parameters and other layers' bias parameters don't use weight decay.
+    If you want to double-check whether the learning rate setup is as expected,
+    you can print the layer names along with their learning rates through
+    "print("Param groups = %s" % json.dumps(parameter_group_names, indent=2))".
+
+    Parameters
+    ----------
+    model
+        A Pytorch model.
+    lr
+        The learning rate.
+    lr_decay
+        The learning rate decay factor (0, 1).
+    weight_decay
+        Weight decay.
+    peft
+        Efficient finetuning strategy. It will only finetune part of the parameters
+    trainable_param_names
+        A list of trainable parameters. (Optional)
+    exclude_keys
+        A list of keys to be excluded.
+
+    Returns
+    -------
+    The grouped parameters based on their layer ids and whether using weight decay.
+    """
+    parameter_group_names = {}
+    parameter_group_vars = {}
+    decay_param_names = get_weight_decay_param_names(model)
+
+    for name, param in model.named_parameters():
+        if name.startswith("_orig_mod."):
+            name = "".join(name.split("_orig_mod."))
+        if exclude_keys and any([exc in name for exc in exclude_keys]):
+            continue
+        layer_id = model.name_to_id[name]
+        if layer_id == 0:  # Set top layer (e.g. head, fusion_mlp, adapter) as being trainable.
+            param.requires_grad = True
+        elif (
+            peft is not None
+            and trainable_param_names
+            and not any([re.match(trainable_param_name, name) for trainable_param_name in trainable_param_names])
+        ):
+            param.requires_grad = False
+
+        if not param.requires_grad:
+            continue  # frozen weights
+
+        if name in decay_param_names:
+            group_name = "decay"
+            this_weight_decay = weight_decay
+        else:
+            group_name = "no_decay"
+            this_weight_decay = 0.0
+
+        layer_id = model.name_to_id[name]
+        group_name = "layer_%d_%s" % (layer_id, group_name)
+
+        if group_name not in parameter_group_names:
+            scale = lr_decay**layer_id
+            parameter_group_names[group_name] = {
+                "weight_decay": this_weight_decay,
+                "params": [],
+                "lr": scale * lr,
+            }
+            parameter_group_vars[group_name] = {
+                "weight_decay": this_weight_decay,
+                "params": [],
+                "lr": scale * lr,
+            }
+
+        parameter_group_vars[group_name]["params"].append(param)
+        parameter_group_names[group_name]["params"].append(name)
+
+    return list(parameter_group_vars.values())

autogluon/multimodal/optim/metrics/__init__.py

@@ -0,0 +1,4 @@
+from .coverage_metrics import Coverage
+from .hit_rate_metrics import CustomHitRate
+from .ranking_metrics import compute_ranking_score
+from .utils import get_torchmetric, compute_score, get_minmax_mode, get_stopping_threshold, infer_metrics

autogluon/multimodal/optim/metrics/coverage_metrics.py

@@ -0,0 +1,42 @@
+import torch
+from torchmetrics import Metric
+from torchmetrics.utilities import dim_zero_cat
+
+
+class Coverage(Metric):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.add_state("pos_probs", default=[], dist_reduce_fx="cat")
+        self.add_state("targets", default=[], dist_reduce_fx="cat")
+        self.tp_threshold = 0.97
+        self.tn_threshold = 0.99
+        higher_is_better = True
+
+    def update(self, preds: torch.Tensor, target: torch.Tensor) -> None:
+        assert preds.dim() == 1
+        assert target.dim() == 1
+        self.pos_probs.append(preds)
+        self.targets.append(target)
+
+    def compute(self):
+        # parse inputs
+        pos_probs = dim_zero_cat(self.pos_probs)
+        targets = dim_zero_cat(self.targets)
+        y_pos = targets[torch.where(pos_probs >= self.tp_threshold)]
+        y_neg = targets[torch.where(pos_probs <= 1 - self.tn_threshold)]
+        tp = sum(y_pos == 1)
+        tn = sum(y_neg == 0)
+        if len(y_pos) == 0 or len(targets) == 0:
+            tp_precision, tp_coverage = 0, 0
+        else:
+            tp_precision, tp_coverage = tp / len(y_pos), len(y_pos) / len(targets)
+            if tp_precision < self.tp_threshold:
+                tp_coverage = 0
+        if len(y_neg) == 0 or len(targets) == 0:
+            tn_precision, tn_coverage = 0, 0
+        else:
+            tn_precision, tn_coverage = tn / len(y_neg), len(y_neg) / len(targets)
+            if tn_precision < self.tn_threshold:
+                tn_coverage = 0
+
+        return torch.tensor(tp_coverage) + torch.tensor(tn_coverage)

autogluon/multimodal/optim/metrics/hit_rate_metrics.py

@@ -0,0 +1,78 @@
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torchmetrics
+
+
+class CustomHitRate(torchmetrics.Metric):
+    """
+    Compute the hit rate when doing semantic search between two group of embeddings.
+    We assume that (a_i, p_i) are a positive pair and (a_i, p_j) for i!=j a negative pair.
+    """
+
+    def __init__(
+        self,
+    ):
+        super().__init__()
+        self.add_state("query_embeddings", default=[], dist_reduce_fx=None)
+        self.add_state("response_embeddings", default=[], dist_reduce_fx=None)
+        self.add_state("logit_scale", default=[], dist_reduce_fx=None)
+
+    def update(
+        self,
+        batch_query_embeds: torch.Tensor,
+        batch_response_embeds: torch.Tensor,
+        logit_scale: Optional[torch.Tensor] = None,
+    ):
+        self.query_embeddings.append(batch_query_embeds)
+        self.response_embeddings.append(batch_response_embeds)
+        if logit_scale is not None:
+            self.logit_scale.append(logit_scale)
+
+    def compute(self):
+        query_embeddings = torch.cat(self.query_embeddings)
+        response_embeddings = torch.cat(self.response_embeddings)
+        if self.logit_scale:
+            logit_scale = torch.mean(torch.stack(self.logit_scale))
+        else:
+            logit_scale = 1
+
+        return self.compute_hit_rate(query_embeddings, response_embeddings, logit_scale)
+
+    @staticmethod
+    def compute_hit_rate(features_a, features_b, logit_scale, top_ks=[1, 5, 10]):
+        """
+        Compute symmetric hit rates between two groups of features.
+
+        Parameters
+        ----------
+        features_a
+            One group of features.
+        features_b
+            The other group of features.
+        logit_scale
+            The scale of logit (Used in CLIP).
+        top_ks
+            Consider only the top k elements for each query.
+
+        Returns
+        -------
+        The accumulated hit rate.
+        """
+        assert len(features_a) == len(features_b)
+        hit_rate = 0
+        logits_per_a = (logit_scale * features_a @ features_b.t()).detach().cpu()
+        logits_per_b = logits_per_a.t().detach().cpu()
+
+        logits = {"logits_per_a": logits_per_a, "logits_per_b": logits_per_b}
+        ground_truth = torch.arange(len(features_b)).view(-1, 1)
+
+        for name, logit in logits.items():
+            ranking = torch.argsort(logit, descending=True)
+            preds = torch.where(ranking == ground_truth)[1]
+
+            for k in top_ks:
+                hit_rate += (preds < k).float().mean()
+
+        hit_rate /= len(top_ks) * len(logits)
+        return hit_rate

autogluon/multimodal/optim/metrics/ranking_metrics.py

@@ -0,0 +1,231 @@
+import logging
+import math
+import operator
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+
+from ...constants import MAP, NDCG, PRECISION, RECALL
+
+logger = logging.getLogger(__name__)
+
+
+class RankingMetrics:
+    def __init__(
+        self,
+        pred: Dict[str, Dict],
+        target: Dict[str, Dict],
+        is_higher_better=True,
+    ):
+        """
+        Evaluation Metrics for information retrieval tasks such as document retrieval, image retrieval, etc.
+        Reference: https://www.cs.cornell.edu/courses/cs4300/2013fa/lectures/metrics-2-4pp.pdf
+
+        Parameters
+        ----------
+        pred:
+            the prediction of the ranking model. It has the following form.
+            pred = {
+                'q1': {
+                    'd1': 1,
+                    'd3': 0,
+                },
+                'q2': {
+                    'd2': 1,
+                    'd3': 1,
+                },
+            }
+            where q refers to queries, and d refers to documents, each query has a few relevant documents.
+            0s and 1s are model predicted scores (does not need to be binary).
+        target:
+            the ground truth query and response relevance which has the same form as pred.
+        is_higher_better:
+            if higher relevance score means higher ranking.
+            if the relevance score is cosine similarity / dot product, it should be set to True;
+            if it is Eulidean distance, it should be False.
+        """
+        self.pred = pred
+        self.target = target
+        self.is_higher_better = is_higher_better
+        # the supported metrics in this script
+        self.supported_metrics = {
+            "precision": 0,
+            "recall": 1,
+            "mrr": 2,
+            "map": 3,
+            "ndcg": 4,
+        }
+
+        assert len(pred) == len(
+            target
+        ), f"The prediction and groudtruth target should have the same number of queries, \
+        while there are {len(pred)} queries in prediction and {len(target)} in the target."
+
+        self.results = {}
+        for key in target.keys():
+            self.results.update({key: [target[key], pred[key]]})
+
+    def compute(self, metrics: Union[str, list] = None, k: Optional[int] = 10):
+        """
+        compute and return ranking scores.
+
+        Parameters
+        ----------
+        metrics:
+            user provided metrics
+        k:
+            the cutoff value for NDCG, MAP, Recall, MRR, and Precision
+
+        Returns
+        -------
+        Computed score.
+
+        """
+        if isinstance(metrics, str):
+            metrics = [metrics]
+        if not metrics:  # no metric is provided
+            metrics = self.supported_metrics.keys()
+
+        return_res = {}
+
+        eval_res = np.mean(
+            [list(self._compute_one(idx, k)) for idx in self.results.keys()],
+            axis=0,
+        )
+
+        for metric in metrics:
+            metric = metric.lower()
+            if metric in self.supported_metrics:
+                return_res.update({f"{metric}@{k}": eval_res[self.supported_metrics[metric]]})
+
+        return return_res
+
+    def _compute_one(self, idx, k):
+        """
+        compute and return the ranking scores for one query.
+        for definition of these metrics, please refer to
+        https://www.cs.cornell.edu/courses/cs4300/2013fa/lectures/metrics-2-4pp.pdf
+
+        Parameters
+        ----------
+        idx:
+            the index of the query
+        k:
+            the cutoff value for NDCG, MAP, Recall, MRR, and Precision
+
+        Returns
+        -------
+        Computed score.
+        """
+        precision, recall, mrr, mAP, ndcg = 0, 0, 0, 0, 0
+        target, pred = self.results[idx][0], self.results[idx][1]
+
+        # sort the ground truth and predictions in descending order
+        sorted_target = dict(
+            sorted(
+                target.items(),
+                key=operator.itemgetter(1),
+                reverse=self.is_higher_better,
+            )
+        )
+        sorted_pred = dict(
+            sorted(
+                pred.items(),
+                key=operator.itemgetter(1),
+                reverse=self.is_higher_better,
+            )
+        )
+        sorted_target_values = list(sorted_target.values())
+        sorted_pred_values = list(sorted_pred.values())
+
+        # number of positive relevance in target
+        # negative numbers and zero are considered as negative response
+        num_pos_target = len([val for val in sorted_target_values if val > 0])
+
+        at_k = k if num_pos_target > k else num_pos_target
+
+        first_k_items_list = list(sorted_pred.items())[0:k]
+
+        rank = 0
+        hit_rank = []  # correctly retrieved items
+        for key, value in first_k_items_list:
+            if key in sorted_target and sorted_target[key] > 0:
+                hit_rank.append(rank)
+            rank += 1
+        count = len(hit_rank)
+        # compute the precision and recall
+        precision = count / k
+        recall = count / num_pos_target
+
+        dcg = 0
+        if hit_rank:  # not empty
+            # compute the mean reciprocal rank
+            mrr = 1 / (hit_rank[0] + 1)
+            # compute the mean average precision
+            mAP = np.sum([sorted_pred_values[rank] * (i + 1) / (rank + 1) for i, rank in enumerate(hit_rank)])
+            # compute the discounted cumulative gain
+            dcg = np.sum([1 / math.log(rank + 2, 2) for rank in hit_rank])
+
+        # compute the ideal discounted cumulative gain
+        idcg = np.sum([1 / math.log(i + 2, 2) for i in range(at_k)])
+        # compute the normalized discounted cumulative gain
+        ndcg = dcg / idcg
+        mAP /= at_k
+
+        return precision, recall, mrr, mAP, ndcg
+
+
+def compute_ranking_score(
+    results: Dict[str, Dict],
+    qrel_dict: Dict[str, Dict],
+    metrics: List[str],
+    cutoffs: Optional[List[int]] = [5, 10, 20],
+):
+    """
+    Compute the ranking metrics, e.g., NDCG, MAP, Recall, and Precision.
+    TODO: Consider MRR.
+
+    Parameters
+    ----------
+    results:
+        The query/document ranking list by the model.
+    qrel_dict:
+        The groundtruth query and document relevance.
+    metrics
+        A list of metrics to compute.
+    cutoffs:
+        The cutoff values for NDCG, MAP, Recall, and Precision.
+
+    Returns
+    -------
+    A dict of metric scores.
+    """
+    scores = {}
+    evaluator = RankingMetrics(pred=results, target=qrel_dict)
+    for k in cutoffs:
+        scores.update(evaluator.compute(k=k))
+
+    metric_results = dict()
+    for k in cutoffs:
+        for per_metric in metrics:
+            if per_metric.lower() == NDCG:
+                metric_results[f"{NDCG}@{k}"] = 0.0
+            elif per_metric.lower() == MAP:
+                metric_results[f"{MAP}@{k}"] = 0.0
+            elif per_metric.lower() == RECALL:
+                metric_results[f"{RECALL}@{k}"] = 0.0
+            elif per_metric.lower() == PRECISION:
+                metric_results[f"{PRECISION}@{k}"] = 0.0
+
+    for k in cutoffs:
+        for per_metric in metrics:
+            if per_metric.lower() == NDCG:
+                metric_results[f"{NDCG}@{k}"] = round(scores[f"{NDCG}@{k}"], 5)
+            elif per_metric.lower() == MAP:
+                metric_results[f"{MAP}@{k}"] = round(scores[f"{MAP}@{k}"], 5)
+            elif per_metric.lower() == RECALL:
+                metric_results[f"{RECALL}@{k}"] = round(scores[f"{RECALL}@{k}"], 5)
+            elif per_metric.lower() == PRECISION:
+                metric_results[f"{PRECISION}@{k}"] = round(scores[f"{PRECISION}@{k}"], 5)
+
+    return metric_results