mAP compute for object detection very slow

[cotrane]

Hi, I'm using torchmetrics.detection.mean_ap.MeanAveragePrecision in finetuning a Detr-Huggingface model and noticed that the compute() call for mAP takes a very long time. I'm training the model on a 4-gpu instance using ddp parallelism.

I'm computing the metric for my validation set of ~6k images which include around ~100k ground-truth objects across 13 classes in them. I do also compute the per class metrics for it. For this setting the compute method seems to take around 40 mins to finish which seems very long, even given my fairly large validation set.

Below you see some pseudo-code of my implementation of the metric computation:

class Detr(pl.LightningModule):
    def __init__(self, lr, lr_backbone, weight_decay, id2label):
        ...
        self.model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", 
                                                            num_labels=len(id2label),
                                                            ignore_mismatched_sizes=True)
        self.val_map = MeanAveragePrecision(
            box_format='cxcywh', class_metrics=True)
        self.id2label = id2label
        
    def common_step(self, batch, batch_idx):
        ...
        outputs = self.model(
            pixel_values=pixel_values,
            pixel_mask=pixel_mask,
            labels=labels
        )

        loss = outputs.loss
        loss_dict = outputs.loss_dict
        probas, labels = outputs.logits.softmax(-1)[:, :, :-1].max(dim=2)
        preds = [
            {
                'boxes': box,
                'scores': proba,
                'labels': label
            } for box, proba, label in zip(outputs.pred_boxes, probas, labels)
        ]
        targets = [
            {
                'boxes': target['boxes'],
                'labels': target['class_labels']
            } for target in batch['labels']
        ]
        return loss, loss_dict, (preds, targets)

    def training_step(self, batch, batch_idx):
        ...
        return loss

    def validation_step(self, batch, batch_idx):
        loss, loss_dict, results_map = self.common_step(batch, batch_idx)
        self.val_map.update(results_map[0], results_map[1]) # preds, targets
        ...
        return loss

    def validation_epoch_end(self, outputs):
        metrics = self.val_map.compute()
        
        mAPs = {"valid_" + k: v for k, v in metrics.items()}            
        mAPs_per_class = mAPs.pop("valid_map_per_class")
        mARs_per_class = mAPs.pop("valid_mar_100_per_class")
        
        self.log_dict({'mem_used': mem_used}, sync_dist=True)
        self.log_dict(mAPs, sync_dist=True)
        self.log_dict(
            {
                f"valid_map_{label}": value
                for label, value in zip(self.id2label.values(), mAPs_per_class)
            },
            sync_dist=True,
        )
        self.log_dict(
            {
                f"valid_mar_100_{label}": value
                for label, value in zip(self.id2label.values(), mARs_per_class)
            },
            sync_dist=True,
        )
        self.val_map.reset()

At this point I'm just curious if that is expected due to my setup or if there is something else going on here. 40 mins compute-time basically doubles the runtime for each epoch of my training run so any help would be much appreciated.

Thanks!

[cotrane]

Using only 3 IOU values instead of range(0.5,1,0.05) reduces to compute time to 12 minutes. That's better but it still feels longer than it should be. If anyone has an idea how to further reduce it I'd be very curious!

[celsofranssa]

Same problem with MRR. See the discussion.

[celsofranssa]

I was able to walk around the performance issues by combining torch metrics with ranx (a Python library of fast ranking evaluation metrics):

import pickle
import torch
from ranx import Qrels, Run, evaluate
from torchmetrics import Metric

class MRRMetric(Metric):
    def __init__(self, params):
        super(MRRMetric, self).__init__(compute_on_step=False)
        self.params=params
        self.run = {} # validation queries and documents
        self.relevance_map = self._load_relevance_map() # the relevance map


    def _load_relevance_map(self):
        with open(f"{self.params.relevance_map.dir}relevance_map.pkl", "rb") as relevances_file:
            data = pickle.load(relevances_file)
        relevance_map = {}
        for text_idx, labels_ids in data.items():
            d = {}
            for label_idx in labels_ids:
                d[f"label_{label_idx}"] = 1.0
            relevance_map[f"text_{text_idx}"] = d
        return relevance_map

    def similarities(self, x1, x2):
        """
        Calculates the cosine similarity matrix for every pair (i, j),
        where i is an embedding from x1 and j is another embedding from x2.

        :param x1: a tensors with shape [batch_size, hidden_size].
        :param x2: a tensors with shape [batch_size, hidden_size].
        :return: the cosine similarity matrix with shape [batch_size, batch_size].
        """
        x1 = x1 / torch.norm(x1, dim=1, p=2, keepdim=True)
        x2 = x2 / torch.norm(x2, dim=1, p=2, keepdim=True)
        return torch.matmul(x1, x2.t())

    def update(self, query_ids, query_rpr, document_ids, document_rpr):
        similarities = self.similarities(query_rpr, document_rpr)
        for row, query_idx in enumerate(query_ids.tolist()):
            if f"query_{query_idx}" not in self.run:
                self.run[f"query_{query_idx}"]={}
            for col, document_idx in enumerate(document_ids.tolist()):
                self.run[f"query_{query_idx}"][f"document_{document_idx}"] = similarities[row][col].item()

    def compute(self):
        return evaluate(
            Qrels({key: value for key, value in self.relevance_map.items() if key in self.run.keys()}),
            Run(self.run),
            ["mrr"] # could be any ranx metric (https://amenra.github.io/ranx/metrics/)
        )

    def reset(self) -> None:
        self.run = {}