Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: normalise in place to reduce memory #82

Merged
merged 2 commits into from
Jan 21, 2025

Conversation

japols
Copy link
Member

@japols japols commented Jan 21, 2025

This PR avoids storing large intermediate results during the normalization by computing it in-place, e.g.

x[..., :] = x[..., :] * self._norm_mul + self._norm_add # old version x.mul_(self._norm_mul).add_(self._norm_add) # new version

Note that the in_place argument still works as expected, only the actual calculation is changed to be done in-place.

For 9km data the old version resulted in mini-peaks during the normalizer:

Version Peak memory (GB, during normalization)
old 32.62
new 17.73

This lines up with storing two additional intermediate results (multiplication, addition) @~8GB batch size in the old version.

@japols japols self-assigned this Jan 21, 2025
@HCookie HCookie changed the title Fix: normalize in place to reduce memory fix:Normalise in place to reduce memory Jan 21, 2025
@HCookie HCookie changed the title fix:Normalise in place to reduce memory fix: normalise in place to reduce memory Jan 21, 2025
@HCookie HCookie added the models label Jan 21, 2025
@japols japols merged commit 40dd1a1 into develop Jan 21, 2025
3 checks passed
@japols japols deleted the models_fix/normalize_inplace branch January 21, 2025 17:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
Projects
Status: Done
Development

Successfully merging this pull request may close these issues.

2 participants