update EMA#670
Merged
yutong-xiang-97 merged 2 commits intomainfrom Mar 30, 2026
Merged
Conversation
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR updates exponential moving average (EMA) updates to use in-place/foreach tensor operations, and routes method implementations to the new internal helpers.
Changes:
- Added
update_ema_tensors(foreach-based) andupdate_momentumhelpers in_torch_helpers.py. - Updated object detection
ModelEMAand DINO/DINOv2/DenseCL codepaths to use the new EMA/momentum helpers. - Added unit tests covering the new EMA/momentum helper behavior.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
src/lightly_train/_torch_helpers.py |
Introduces foreach-based EMA tensor updates and a momentum helper. |
src/lightly_train/_task_models/object_detection_components/ema.py |
Switches ModelEMA update logic to batch-update floating tensors via update_ema_tensors. |
src/lightly_train/_methods/dinov2/dinov2_loss.py |
Uses in-place ops for center EMA update to avoid tensor reassignment. |
src/lightly_train/_methods/dinov2/dinov2.py |
Replaces external update_momentum import with internal helper. |
src/lightly_train/_methods/dino/dino.py |
Replaces external update_momentum import with internal helper. |
src/lightly_train/_methods/densecl/densecl.py |
Uses internal update_momentum helper for key/query encoder momentum update. |
tests/test__torch_helpers.py |
Adds unit tests for update_momentum and dtype-grouped update_ema_tensors. |
yutong-xiang-97
force-pushed
the
yutong-trn-1834-update-momentum-functions
branch
from
e648324 to
413f3eb
Compare
Contributor
Author
|
/review |
guarin
suggested changes
Mar 27, 2026
guarin
approved these changes
Mar 30, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What has changed and why?
Update EMA calculation with
_foreach_mul_How has it been tested?
Unit tests
Did you update CHANGELOG.md?
Did you update the documentation?