Skip to content

Reduce EoMT instance validation mask memory - #942

Open
KrzysiekSko wants to merge 4 commits into
lightly-ai:mainfrom
KrzysiekSko:fix/941-eomt-instance-val-memory
Open

Reduce EoMT instance validation mask memory#942
KrzysiekSko wants to merge 4 commits into
lightly-ai:mainfrom
KrzysiekSko:fix/941-eomt-instance-val-memory

Conversation

@KrzysiekSko

Copy link
Copy Markdown

What has changed and why?

Implements Task 1 from #941 for the DINOv2 and DINOv3 EoMT instance-segmentation validation paths.

  • Processes the query dimension in chunks sized from a 512 MiB estimated working-set budget rather than a fixed query count.
  • Accounts conservatively for six FP32-sized mask-score buffers per output pixel.
  • Resizes and scores one chunk at a time.
  • Moves completed prediction masks to CPU immediately while keeping labels and scores on their original device.
  • Leaves get_labels_masks_scores unchanged.
  • Adds focused tests for byte-budget chunk sizing, actual chunk boundaries, and numerical equality with the unchunked path.

The metric semantics and original-image evaluation resolution are unchanged.

Refs #941

How has it been tested?

The focused tests are included at:

tests/_task_models/eomt/test_instance_segmentation.py

They compare all labels, masks, and scores from a forced multi-chunk run against a single-chunk run using deterministic tensors.

A full local test run and GPU peak-memory measurement were not available in this environment. CI is requested for the repository test matrix. The device check remaining for review is validation on the same checkpoint/data before and after the change, including torch.cuda.max_memory_allocated() on a wide image.

Did you update CHANGELOG.md?

  • Yes
  • Not needed (internal change)

Did you update the documentation?

  • Yes
  • Not needed (internal change without effects for user)

@CLAassistant

CLAassistant commented Aug 12, 2026

Copy link
Copy Markdown

CLA assistant check
All committers have signed the CLA.

@KrzysiekSko
KrzysiekSko marked this pull request as ready for review August 12, 2026 13:18

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 5570d08ac1

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

Comment thread src/lightly_train/_task_models/eomt/instance_segmentation.py Outdated
@KrzysiekSko

Copy link
Copy Markdown
Author

Addressed in c21701f and formatted in 6015441.

The chunk budget now uses the larger materialized mask surface by pixel count across model_image_size and the final image_size, so the first resize cannot bypass the budget when the final image is smaller. I added a regression where model_image_size=(8, 10) and image_size=(4, 5); a two-query model-resolution budget produces chunk boundaries [2, 2, 1].

Validation:

  • uv run pytest tests/_task_models/eomt/test_instance_segmentation.py -q: 3 passed
  • focused Ruff check: passed
  • focused Ruff format check: passed
  • git diff --check: passed

@liopeer

liopeer commented Aug 12, 2026

Copy link
Copy Markdown
Contributor

/review

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants