lightly-ai / lightly-ai/lightly-train

EoMT validation uses much more GPU memory than training

Open
#941 3 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug dinov2 dinov3 help wanted segmentation tracking issue
Dominant language
Python
Stars
1.7k
Forks
116
Avg merge
2d 21h
Merged PRs (30d)
6

Description

## Problem

EoMT models run out of GPU memory in the validation step. Training with the same batch size works.

Issue #936 reports this. A user trains `dinov3/vitt16plus-eomt-inst-coco` with `batch_size=1` on a 12 GB card. Training uses about 2.5 GB. Validation stops with `torch.OutOfMemoryError: Tried to allocate 3.15 GiB`.

We want to lower the memory use. The metric values must stay the same.

## Cause

The validation step resizes the predicted masks to the original image size. This follows the EoMT evaluation protocol. The step resizes all queries together.

Here are the numbers from #936. The model has 200 queries. The original image is 2943x1435 pixels. One float32 tensor of shape `(1, 200, 1435, 2943)` needs 3.15 GiB. This is the exact number in the error message.

`get_labels_masks_scores` then makes four more tensors of the same shape: `masks`, `masks_fp32`, the output of `sigmoid()`, and their product. The peak is near 13 GiB.

Training does not have this problem. Training resizes the masks to the model input size, for example 640x640. That tensor is 10 times smaller. Training also skips this code by default, because `metric_args.train` is `False`.

## Tasks

Please start with task 1. We want to agree on the approach on a small diff first. Task 2 and task 3 can follow in separate pull requests.

### Task 1 — instance segmentation (start here)

Files:

- `src/lightly_train/_task_models/dinov2_eomt_instance_segmentation/train_model.py:403-427`
- `src/lightly_train/_task_models/dinov3_eomt_instance_segmentation/train_model.py:407-434`

Split the query dimension into chunks. For each chunk, resize the mask logits, then call `get_labels_masks_scores`. Keep the results and free the chunk.

This gives the same numbers. Bilinear interpolation works on each channel alone. The score reduction `.sum(2)` works on each query alone. A chunk boundary changes neither.

Pick the chunk size from a byte budget, not from a fixed count. Large images then get small chunks. `linear_semantic_segmentation/train_model.py:212-235` shows the same idea for tiles.

Please do not change `get_labels_masks_scores` in `task_model.py`. That code carries an FP32 and TensorRT warning, and the ONNX export path uses it.

Two more steps are free, and we checked that both keep the results:

- Move the predicted masks to the CPU before the metric update. torchmetrics already calls `.cpu()` on them. The plot code also calls `.cpu()` on them.
- Keep `labels` and `scores` on the GPU. torchmetrics syncs those tensors across ranks, and NCCL needs CUDA tensors.

### Task 2 — panoptic segmentation

Files:

- `src/lightly_train/_task_models/dinov2_eomt_panoptic_segmentation/train_model.py:462-468`
- `src/lightly_train/_task_models/dinov3_eomt_panoptic_segmentation/train_model.py:466-472`

Panoptic needs a different fix. `get_image_masks_segment_ids_scores` runs `argmax` over all queries for each pixel. Query chunks do not work there.

But the function drops queries first:

```python
keep = (labels != ignore_class_id) & (scores > threshold)
```

This test reads `class_logits` only. It does not read the masks. So we can drop the queries before the resize, not after. The kept queries then get resized alone.

This gives the same numbers, because interpolation works on each channel alone.

### Task 3 — semantic segmentation

Files:

- `src/lightly_train/_task_models/dinov2_eomt_semantic_segmentation/train_model.py:357-358`
- `src/lightly_train/_task_models/dinov3_eomt_semantic_segmentation/train_model.py:359-360`

This is a different problem. The step cuts each image into tiles and stacks all tiles into one forward pass:

```python
crops_list, origins = self.model.tile(images)
crops = torch.stack(crops_list)
```

The tile count grows with the image aspect ratio. A 4:1 image gives 4 tiles. So the forward batch is larger than the dataloader batch, and the size is not bounded.

`linear_semantic_segmentation` fixed the same pattern in PR #880. Please copy that approach: run the tiles in chunks, and add each chunk to the accumulators at once.

## How to check a fix

1. Run validation before and after your change on the same checkpoint and the same data. All metric values must match.
2. Read the peak memory with `torch.cuda.max_memory_allocated()`. Compare the training phase and the validation phase.
3. Use a wide image to see the effect. An image near 3000x1500 pixels works well.
4. Add a test next to the existing tests, for example `tests/_task_models/dinov3_eomt_instance_segmentation/test_train_model.py`.

## Out of scope

Two points are known, and this issue does not cover them.

The metric keeps every prediction until the end of the validation phase. `MeanAveragePrecision` stores all 200 query masks for each image. A top-N filter can cut this, but pycocotools trims per class, so a global top-N gives different numbers. We will track this apart.

#936 also asks for an option to compute the metrics at the transform size. That option changes the metric values, so it does not belong here.

## How to help

Comment on this issue to claim a task. Ask any question here. We are happy to help with the numerical checks.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the validation code in src/lightly_train/_task_models/dinov2_eomt_instance_segmentation/train_model.py and the corresponding DINOv3 file at the listed lines. Compare the chunking approach in linear_semantic_segmentation/train_model.py, then run tests/_task_models/dinov3_eomt_instance_segmentation/test_train_model.py and validation on a wide image. Done means lower peak GPU memory with unchanged metric values.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
computer-vision, performance
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
72/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.