search

DiffEqML / torchdyn

A PyTorch library entirely dedicated to neural differential equations, implicit models and related numerical methods
https://torchdyn.org
Apache License 2.0
1.35k stars 125 forks source link

How to get memory usage for "adjoint" and "autograd" method? #138

Open cyx96 opened 2 years ago

cyx96 commented 2 years ago

Thanks for this amazing package!

I was trying to test the memory usage of adjoint, as claimed by authors of the original neural ODE paper, the memory usage of adjoint method should be smaller compared to vanilla "autograd". However, the output of torch.cuda.memory_summary() show an increase of GPU memory of the adjoint method compared to autograd. I'm wondering if I used torch.cuda.memory_summary() wrong, I printed it after the training. If my approach was incorrect, what is the correct way to get memory usage for "adjoint" and "autograd" method?

joglekara commented 2 years ago

Hey, did you happen to make progress on this? I am curious to know and can hopefully provide some benchmarks when I get my problem running as well