Memory profiling

[wiederm]

Pull Request Summary

This PR adds infrastructure to perform GPU memory profiling as described in this series of blog posts by the PyTorch folks. I have added functions to initialize, record and save GPU memory traces and a notebook that uses these functions to profile forward/backward pass on a 25A waterbox.

I also added a notebook that compares the timing and memory GPU consumption of each of the implemented potentials.

This PR also adds optimization for each of the implemented networks.

ANI architecture

Investigating the memory trace shows that the computation of the angular aev allocates the largest junk of the memory:

This is due to the creation of large intermediate tensors due to broadcasting over multiple dimensions; refactoring this saves around 100 MB of GPU memory.

• refactor large memory allocation in angular aev

Key changes

• adding profile module with helper functions to profile GPU memory traces

Associated Issue(s)

• Issue 1

Pull Request Checklist

• Issue(s) raised/addressed and linked
• Includes appropriate unit test(s)
• Appropriate docstring(s) added/updated
• Appropriate .rst doc file(s) added/updated
• PR is ready for review

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 31.57895% with 52 lines in your changes missing coverage. Please review.

Project coverage is 84.53%. Comparing base (80e4f43) to head (44f32f2).

Additional details and impacted files