Clear list states (i.e. delete their contents), not reassign the default []

[dominicgkerr]

Rather than overwriting list states (with []), call .clear() to correctly (more robustly?) delete/free Tensor elements. Previous behaviour results in CPU (possibly GPU also) memory leak

What does this PR do?

Fixes #2492 (maybe #2481 also)

Before submitting

• Was this discussed via a Github issue? (no need for typos and docs improvements)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure to update the docs?
• Did you write any new necessary tests?

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

Did you have fun?

Make sure you had fun coding 🙃

---

📚 Documentation preview 📚: https://torchmetrics--2493.org.readthedocs.build/en/2493/

[Borda]

nice, can we also have a test to cover this edge-case?

[dominicgkerr]

nice, can we also have a test to cover this edge-case?

Absolutely! Added 1fa7077 to help illustrate the issue (and how the fix helps)

[justusschock]

Thanks for your PR and that's certainly a great find.
That being said, I am actually not sure what's the desired behavior here.

I can see good reasons for both ways.

I.e. It might be very confusing for people if they keep a reference to things why this one suddenly empty, whereas with just instantiating a new list we're doing everything correctly on our side and still give the user the possibility to retain a reference to the old state if they wish to. So this would only lead to a memory leak in user code, not on our side.

[dominicgkerr]

Interesting - I hadn't thought of that scenario being desirable...

I would argue though, the exiting behaviour is certainly unexpected (according to the current documentation). My reading of .add_state is that users (intentionally) defer the memory/device management entirely to Metric, and wouldn't use it if they needed something more complex (like persistent references.) Similarly, I naturally expect .reset to return the Metric to its original states, AND to tidy-up any intermediate values calculated by .update (which are currently leaked.)

In my specific use-case, I'm fitting a LighntningModule to a very large dataset. I have a setup very similar to [1], except I'm using a custom Metric with a list state - over the course of a large number of training steps, I consistently see my CPU memory usage increase (as a result of #2492) until the python process crashes.

If you're keen to support the ability to retain references to Metric states, would introducing a (say) free_on_reset: bool keyword argument to .add_state (that conditionally deletes/clears states before reassigning the default value) be preferable? Setting free_on_reset=True by could enable the new behaviour by default, while still allowing users to revert to the existing behaviour whn they require something more complex?

• [1] https://lightning.ai/docs/torchmetrics/stable/pages/lightning.html#logging-torchmetrics

[justusschock]

I think you have some valid arguments there. I'm fine making an opinionated move here. Could you just add that to the documentation and refer people to use copy/deepcopy if the want to retain the states?

[dominicgkerr]

Brilliant, thanks! I've updated the docstring/documentation - let me know if it needs rewording/expanding/etc...

Unfortunately I couldn't get make docs to build locally (as I couldn't pip install lai-sphinx-theme), so just waiting on the CI to double check my formatting is correct

Thanks again

[Borda]

@SkafteNicki thoughts?

[Borda]

@dominicgkerr there are too many failing tests, could you pls have look...

[dominicgkerr]

@SkafteNicki Nice (5758977) - was just looking at the .forward logic!

[SkafteNicki]

@dominicgkerr you are welcome :)
Hopefully all tests should pass now

[codecov]

Codecov Report

Merging #2493 (3f013cf) into master (581c444) will increase coverage by 0%.
The diff coverage is 100%.

Additional details and impacted files

@@          Coverage Diff           @@
##           master   #2493   +/-   ##
======================================
  Coverage      69%     69%           
======================================
  Files         307     307           
  Lines       17396   17404    +8     
======================================
+ Hits        11981   11989    +8     
  Misses       5415    5415           

[dominicgkerr]

I think I've got things working (better?) now... @SkafteNicki / 5758977 was a huge help (thanks!), as I hadn't spotted the caching behaviour inside .forward / .update before.

Rather than simply using deepcopy (which worked in all but one edge case), I added a ._copy_state_dict helper method to .detach().clone() Tensor / list[Tensor] states inserted into the cache. Unfortunately this caused tests/unittests/bases/test_metric.py:test_constant_memory_on_repeat_init to fail, but I found (in ef27215) that removing the .clone() calls, fixed it. But, as the resulting cache might now contain references to Tensor values, I was suspicious that it somehow leak memory...

After some digging, I came to the conclusion that test_constant_memory_on_repeat_init wasn't actually testing memory leakage during initialisation, but actually ensure repeated .forward calls didn't (significantly) increase the metric's memory. As an illustration, moving x = torch.randn(10000).cuda() inside the test's for loop (where a constant Tensor couldn't be automatically referenced by python/pytorch in subsequent calls), actually caused it to fail...

I claim, memory should be allowed to increase here (as users might legitimately want to collect lots of observations during fitting, and combine them inside .compute), as long as the allocated memory is eventually freed by .reset. I've added a new test (test_freed_memory_on_reset) check this.

I appreciate changing (failing) tests to pass the CI is pretty suspect - hopefully the above explains why I did (let me know if not!)