lightly-ai / lightly-ai/lightly-train
EoMT validation uses much more GPU memory than training
Nobody has claimed this yet.
- 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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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