CUDA Out Of Memory when training a DETR Object detection model with compute_metrics
Author: Kamal-Moha
Pull request: no
Labels: name: trainer

name: bug

name: Vision

### System Info

`transformers`  version 4.47.0.dev0
`accelerate`  version 0.34.2
`timm`  version 1.0.11
`supervision`  version 0.25.0rc2

### Who can help?

@muellerzr @Arth

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

I’m training a DETR Object Detection model using the Trainer API. I have properly created the coco dataset.

But when I run the Trainer API with `custom_metrics`, I get the error “OutOfMemoryError: CUDA out of memory.”. I have reduced batch_size from 16 until 1, but the same error of “Out of memory”. 

```
def collate_fn(batch):
    data = {}
    data["pixel_values"] = torch.stack([x["pixel_values"] for x in batch])
    data["labels"] = [x["labels"] for x in batch]
    return data
```

Here’s how I’m creating the `custom_metrics` function

```
id2label = {id: label for id, label in enumerate(train_ds.classes)}
label2id = {label: id for id, label in enumerate(train_ds.classes)}


@dataclass
class ModelOutput:
    logits: torch.Tensor
    pred_boxes: torch.Tensor

class MAPEvaluator:

    def __init__(self, image_processor, threshold=0.00, id2label=None):
        self.image_processor = image_processor
        self.threshold = threshold
        self.id2label = id2label

    def collect_image_sizes(self, targets):
        """Collect image sizes across the dataset as list of tensors with shape [batch_size, 2]."""
        image_sizes = []
        for batch in targets:
            batch_image_sizes = torch.tensor(np.array([x["size"] for x in batch]))
            image_sizes.append(batch_image_sizes)
        return image_sizes

    def collect_targets(self, targets, image_sizes):
        post_processed_targets = []
        for target_batch, image_size_batch in zip(targets, image_sizes):
            for target, (height, width) in zip(target_batch, image_size_batch):
                boxes = target["boxes"]
                boxes = sv.xcycwh_to_xyxy(boxes)
                boxes = boxes * np.array([width, height, width, height])
                boxes = torch.tensor(boxes)
                labels = torch.tensor(target["class_labels"])
                post_processed_targets.append({"boxes": boxes, "labels": labels})
        return post_processed_targets

    def collect_predictions(self, predictions, image_sizes):
        post_processed_predictions = []
        for batch, target_sizes in zip(predictions, image_sizes):
            batch_logits, batch_boxes = batch[1], batch[2]
            output = ModelOutput(logits=torch.tensor(batch_logits), pred_boxes=torch.tensor(batch_boxes))
            post_processed_output = self.image_processor.post_process_object_detection(
                output, threshold=self.threshold, target_sizes=target_sizes
            )
            post_processed_predictions.extend(post_processed_output)
        return post_processed_predictions

    @torch.no_grad()
    def __call__(self, evaluation_results):

        predictions, targets = evaluation_results.predictions, evaluation_results.label_ids

        image_sizes = self.collect_image_sizes(targets)
        post_processed_targets = self.collect_targets(targets, image_sizes)
        post_processed_predictions = self.collect_predictions(predictions, image_sizes)

        evaluator = MeanAveragePrecision(box_format="xyxy", class_metrics=True)
        evaluator.warn_on_many_detections = False
        evaluator.update(post_processed_predictions, post_processed_targets)

        metrics = evaluator.compute()

        # Replace list of per class metrics with separate metric for each class
        classes = metrics.pop("classes")
        map_per_class = metrics.pop("map_per_class")
        mar_100_per_class = metrics.pop("mar_100_per_class")
        for class_id, class_map, class_mar in zip(classes, map_per_class, mar_100_per_class):
            class_name = id2label[class_id.item()] if id2label is not None else class_id.item()
            metrics[f"map_{class_name}"] = class_map
            metrics[f"mar_100_{class_name}"] = class_mar

        metrics = {k: round(v.item(), 4) for k, v in metrics.items()}

        return metrics

eval_compute_metrics_fn = MAPEvaluator(image_processor=processor, threshold=0.01, id2label=id2label)
```
This is how I'm training the model

```
training_args = TrainingArguments(
    output_dir=f"Malaria-finetune",
    report_to="none",
    num_train_epochs=10,
    max_grad_norm=0.1,
    learning_rate=5e-5,
    warmup_steps=300,
    per_device_train_batch_size=1,
    dataloader_num_workers=2,
    metric_for_best_model="eval_map",
    greater_is_better=True,
    load_best_model_at_end=True,
    eval_strategy="epoch",
    save_strategy="epoch",
    save_total_limit=2,
    remove_unused_columns=False,
    eval_do_concat_batches=False,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=pytorch_dataset_train,
    eval_dataset=pytorch_dataset_valid,
    processing_class=processor,
    data_collator=collate_fn,
    compute_metrics=eval_compute_metrics_fn
)
trainer.train()

```

I would like to please get help on this. 

### Expected behavior

I expected the Trainer to train normally on this custom dataset for 10 epochs without any errors

Comments:
Hi, what GPU are you using? 


@Haroon-Mughal I'm using the Kaggle GPU P100 and also the colab GPU for this

cc @qubvel 

Hi @Kamal-Moha, thanks for the issue!
You might have a huge evaluation dataset and/or strong memory constraints, thus storing all the validation dataset bounding boxes might lead to an OutOfMemory error. Can you confirm that avoiding passing `compute_metrics` allows you to run the script without any issues? If so there is a way to implement per-batch metrics computation similar to instance-segmentation example (see [here](https://github.com/huggingface/transformers/blob/13493215abceafc1653af88b045120014fb4c1fc/examples/pytorch/instance-segmentation/run_instance_segmentation.py#L171) and [here](https://github.com/huggingface/transformers/blob/13493215abceafc1653af88b045120014fb4c1fc/examples/pytorch/instance-segmentation/run_instance_segmentation.py#L356)). Please let me know if that helps or you have more questions!

Hi @qubvel, thanks for your reply!

YES I confirm that when I don't pass `compute_metrics`, my training goes on well without any issues. OutOfMemory error occurs only when I pass `compute_metrics`. Would like to know how to implement the per-batch metrics computation for object detection.

Hi @Kamal-Moha, the general idea is to set the [training argument](https://huggingface.co/docs/transformers/main/main_classes/trainer#transformers.Seq2SeqTrainingArguments.batch_eval_metrics) `batch_eval_metrics=True` and implement the function/class method for computing metrics on each batch, see the example I linked above. You can reuse `Evaluator` class from instance-segmentation example, just make sure to update it with boxed instead of masks.
Please, let me know if you can adapt the code for object detection 🤗 



@qubvel Could you please help me implement the class method of computing metrics for object detection. 

Below is my attempt to adapt the class for object detection;
```
class Evaluator:
    """
    Compute metrics for the instance segmentation task.
    """

    def __init__(
        self,
        image_processor: AutoImageProcessor,
        id2label: Mapping[int, str],
        threshold: float = 0.0,
    ):
        """
        Initialize evaluator with image processor, id2label mapping and threshold for filtering predictions.

        Args:
            image_processor (AutoImageProcessor): Image processor for
                `post_process_instance_segmentation` method.
            id2label (Mapping[int, str]): Mapping from class id to class name.
            threshold (float): Threshold to filter predicted boxes by confidence. Defaults to 0.0.
        """
        self.image_processor = image_processor
        self.id2label = id2label
        self.threshold = threshold
        self.metric = self.get_metric()

    def get_metric(self):
        metric = MeanAveragePrecision(iou_type="bbox", class_metrics=True)
        return metric

    def reset_metric(self):
        self.metric.reset()

    def postprocess_target_batch(self, target_batch) -> List[Dict[str, torch.Tensor]]:
        """Collect targets in a form of list of dictionaries with keys "masks", "labels"."""
        batch_boxes = target_batch[0]
        batch_labels = target_batch[1]
        post_processed_targets = []
        for boxes, labels in zip(batch_boxes, batch_labels):
            post_processed_targets.append(
                {
                    # "boxes": boxes.to(dtype=torch.bool),
                    "boxes": boxes,
                    "labels": labels,
                }
            )
        return post_processed_targets

    def get_target_sizes(self, post_processed_targets) -> List[List[int]]:
        target_sizes = []
        for target in post_processed_targets:
            target_sizes.append(target["boxes"].shape[-2:])
        return target_sizes

    def postprocess_prediction_batch(self, prediction_batch, target_sizes) -> List[Dict[str, torch.Tensor]]:
        """Collect predictions in a form of list of dictionaries with keys "masks", "labels", "scores"."""

        model_output = ModelOutput(class_queries_logits=prediction_batch[0], masks_queries_logits=prediction_batch[1])
        post_processed_output = self.image_processor.post_process_instance_segmentation(
            model_output,
            threshold=self.threshold,
            target_sizes=target_sizes,
            return_binary_maps=True,
        )

        post_processed_predictions = []
        for image_predictions, target_size in zip(post_processed_output, target_sizes):
            if image_predictions["boxes"]:
                post_processed_image_prediction = {
                    "boxes": image_predictions["boxes"].to(dtype=torch.bool),
                    "labels": torch.tensor([x["label_id"] for x in image_predictions["boxes"]]),
                    "scores": torch.tensor([x["score"] for x in image_predictions["boxes"]]),
                }
            else:
                # for void predictions, we need to provide empty tensors
                post_processed_image_prediction = {
                    "boxes": torch.zeros([0, *target_size], dtype=torch.bool),
                    "labels": torch.tensor([]),
                    "scores": torch.tensor([]),
                }
            post_processed_predictions.append(post_processed_image_prediction)

        return post_processed_predictions

    @torch.no_grad()
    def __call__(self, evaluation_results: EvalPrediction, compute_result: bool = False) -> Mapping[str, float]:
        """
        Update metrics with current evaluation results and return metrics if `compute_result` is True.

        Args:
            evaluation_results (EvalPrediction): Predictions and targets from evaluation.
            compute_result (bool): Whether to compute and return metrics.

        Returns:
            Mapping[str, float]: Metrics in a form of dictionary {<metric_name>: <metric_value>}
        """
        prediction_batch = nested_cpu(evaluation_results.predictions)
        target_batch = nested_cpu(evaluation_results.label_ids)

        # For metric computation we need to provide:
        #  - targets in a form of list of dictionaries with keys "masks", "labels"
        #  - predictions in a form of list of dictionaries with keys "masks", "labels", "scores"
        post_processed_targets = self.postprocess_target_batch(target_batch)
        target_sizes = self.get_target_sizes(post_processed_targets)
        post_processed_predictions = self.postprocess_prediction_batch(prediction_batch, target_sizes)

        # Compute metrics
        self.metric.update(post_processed_predictions, post_processed_targets)

        if not compute_result:
            return

        metrics = self.metric.compute()

        # Replace list of per class metrics with separate metric for each class
        classes = metrics.pop("classes")
        map_per_class = metrics.pop("map_per_class")
        mar_100_per_class = metrics.pop("mar_100_per_class")
        for class_id, class_map, class_mar in zip(classes, map_per_class, mar_100_per_class):
            class_name = self.id2label[class_id.item()] if self.id2label is not None else class_id.item()
            metrics[f"map_{class_name}"] = class_map
            metrics[f"mar_100_{class_name}"] = class_mar

        metrics = {k: round(v.item(), 4) for k, v in metrics.items()}

        # Reset metric for next evaluation
        self.reset_metric()

        return metrics
```
I mostly just replaced "segmentation" from original Evaluator class(for segmentation) with "boxes"(to see if it works for object detection).


Here's my Trainer;
```
training_args = TrainingArguments(
    output_dir=f"Ghana Crop Detection",
    report_to="none",
    num_train_epochs=10,
    max_grad_norm=0.1,
    learning_rate=5e-5,
    warmup_steps=300,
    per_device_train_batch_size=2,
    per_device_eval_batch_size=2,
    dataloader_num_workers=2,
    metric_for_best_model="eval_map",
    greater_is_better=True,
    load_best_model_at_end=True,
    eval_strategy="epoch",
    save_strategy="epoch",
    save_total_limit=2,
    remove_unused_columns=False,
    eval_do_concat_batches=False,
    batch_eval_metrics=True
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=pytorch_dataset_train,
    eval_dataset=pytorch_dataset_valid,
    processing_class=processor,
    data_collator=collate_fn,
    compute_metrics=compute_metrics
)
trainer.train()
```


When I run it, I get the below error

```
[<ipython-input-48-1f4dd0bb1f9c>](https://localhost:8080/#) in get_target_sizes(self, post_processed_targets)
     49         target_sizes = []
     50         for target in post_processed_targets:
---> 51             target_sizes.append(target["boxes"].shape[-2:])
     52         return target_sizes
     53 

AttributeError: 'str' object has no attribute 'shape' 
```

I need further help, please.


Hey @Kamal-Moha! Please have a look at this PR

- https://github.com/huggingface/transformers/pull/35040

@qubvel I have implemented the code changes from the PR. This is my new Evaluator class

```
class Evaluator:
    # compute metrics for object detection

  def __init__(
        self,
        image_processor: AutoImageProcessor,
        id2label: Mapping[int, str],
        threshold: float = 0.0,
    ):
    self.image_processor = image_processor
    self.id2label = id2label
    self.threshold = threshold
    self.metric = self.get_metric()

  def get_metric(self):
    metric = MeanAveragePrecision(iou_type="bbox", box_format="xyxy", class_metrics=True)
    return metric

  def reset_metric(self):
    self.metric.reset()

  def postprocess_target_batch(self, target_batch) -> List[Dict[str, torch.Tensor]]:
    """Collect targets in a form of list of dictionaries with keys "boxes", "labels"."""
    post_processed_targets = []
    for sample in target_batch:
      boxes = convert_bbox_yolo_to_pascal(sample["boxes"], sample["orig_size"])
      labels = sample["class_labels"]
      post_processed_targets.append({"boxes": boxes, "labels": labels})
    return post_processed_targets

  def get_target_sizes(self, target_batch) -> torch.Tensor:
    target_sizes = torch.stack([x["orig_size"] for x in target_batch])
    return target_sizes

  def postprocess_prediction_batch(self, prediction_batch, target_sizes) -> List[Dict[str, torch.Tensor]]:
    """Collect predictions in a form of list of dictionaries with keys "boxes", "labels", "scores"."""
    batch_logits, batch_boxes = prediction_batch[1], prediction_batch[2]
    output = ModelOutput(logits=batch_logits, pred_boxes=batch_boxes)
    post_processed_predictions = self.image_processor.post_process_object_detection(
        output, threshold=self.threshold, target_sizes=target_sizes
    )
    return post_processed_predictions

    # Compute metrics
  @torch.no_grad()
  def __call__(self, evaluation_results: EvalPrediction, compute_result: bool = False) -> Mapping[str, float]:
    prediction_batch = evaluation_results.predictions
    target_batch = evaluation_results.label_ids
    target_sizes = self.get_target_sizes(target_batch)
    post_processed_targets = self.postprocess_target_batch(target_batch)
    post_processed_predictions = self.postprocess_prediction_batch(prediction_batch, target_sizes)

# move tensors to CPU
    post_processed_targets = nested_cpu(post_processed_targets)
    post_processed_predictions = nested_cpu(post_processed_predictions)

    # Compute metrics
    self.metric.update(post_processed_predictions, post_processed_targets)
    if not compute_result:
        return

    metrics = self.metric.compute()
    # Replace list of per class metrics with separate metric for each class
    classes = metrics.pop("classes")
    map_per_class = metrics.pop("map_per_class")
    mar_100_per_class = metrics.pop("mar_100_per_class")
    for class_id, class_map, class_mar in zip(classes, map_per_class, mar_100_per_class):
        class_name = self.id2label[class_id.item()] if self.id2label is not None else class_id.item()
        metrics[f"map_{class_name}"] = class_map
        metrics[f"mar_100_{class_name}"] = class_mar

    metrics = {k: round(v.item(), 4) for k, v in metrics.items()}
    # Reset metric for next evaluation
    self.reset_metric()

    return metrics
```

Trainer successfully runs for 1 epoch

![image](https://github.com/user-attachments/assets/5591b05e-fa52-404d-9fdf-d4f552d46f78)

But after 1 epoch, I get the below error.
![image](https://github.com/user-attachments/assets/1f57511e-2081-4feb-9a8e-41d8ce9d2325)

I don't know why after 1 epoch, I get this error with the `convert_bbox_yolo_to_pascal` function. I have properly unpacked the image_size to get `height` and `width` of the image. For debugging purposes, I have printed the `height` & `width` when running the Trainer. I can confirm that in the first epoch, height and width have been created properly. 
But when proceeding to another epoch, the error comes up. Don't know why.

Here is my `convert_bbox_yolo_to_pascal` function;

```
def convert_bbox_yolo_to_pascal(boxes: torch.Tensor, image_size: Tuple[int, int]) -> torch.Tensor:
  # convert center to corners format
  boxes = center_to_corners_format(boxes)

  # convert to absolute coordinates
  height, width = image_size[:2]
  print(f"height: {height}, width: {width}")
  boxes = boxes * torch.tensor([[width, height, width, height]], device=boxes.device)
  print(f"boxes: {boxes}")

  return boxes
```

Keen to get your thoughts & help on this. Thanks!

Hey @Haroon-Mughal, I would be glad to help, but unfortunately, I don't have any guess regarding the problem. My approach to debugging would be to log as much as possible to identify the cause. Try running on a smaller (reduced) dataset to check if there is any issue in the data or code, and log shapes, keys, and values. Check if the error occurs at the beginning or on a particular batch. Hope that helps somehow!

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.