fix(diff/implicit): fix memory leak of OOP APIs

[XuehaiPan]

No description provided.

[codecov-commenter]

Codecov Report

Base: 70.34% // Head: 70.26% // Decreases project coverage by -0.07% ⚠️

Coverage data is based on head (d233cbe) compared to base (381f4d1).
Patch coverage: 83.00% of modified lines in pull request are covered.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #113      +/-   ##
==========================================
- Coverage   70.34%   70.26%   -0.08%     
==========================================
  Files          71       72       +1     
  Lines        2981     3000      +19     
==========================================
+ Hits         2097     2108      +11     
- Misses        884      892       +8     

Flag	Coverage Δ
unittests	70.26% <83.00%> (-0.08%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
torchopt/nn/module.py	40.09% <27.27%> (+0.27%)	⬆️
torchopt/optim/meta/base.py	63.15% <75.00%> (ø)
torchopt/nn/stateless.py	81.57% <81.57%> (ø)
torchopt/diff/implicit/nn/module.py	81.39% <96.87%> (-4.32%)	⬇️
torchopt/__init__.py	100.00% <100.00%> (ø)
torchopt/diff/implicit/decorator.py	95.75% <100.00%> (+0.02%)	⬆️
torchopt/diff/zero_order/nn/module.py	91.42% <100.00%> (-1.08%)	⬇️
torchopt/nn/__init__.py	100.00% <100.00%> (ø)
torchopt/typing.py	95.65% <100.00%> (+0.19%)	⬆️
torchopt/utils.py	56.85% <100.00%> (ø)
... and 1 more

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

☔ View full report at Codecov.
📢 Do you have feedback about the report comment? Let us know in this issue.

[zaccharieramzi]

Hi @XuehaiPan @JieRen98 ,

I see that this was just merged and I seem to be running in a memory leak of my own.
Could you explain what the problem was and how it was solved?

EDIT

Do you also know when the release will happen for this fix?

EDIT 2

Apparently my memory leak was not fixed by using the latest fix so my guess is that it's not the same issue.

[XuehaiPan]

Hi @zaccharieramzi, you could reference the changes in examples/iMAML/imaml_omniglot.py (5cb7b6).

We previously created the inner network in the for-loop. For each iteration, a new inner network was created. That was the cause of the memory leak.

for task_id in range(num_tasks):
    inner_net = InnerNet(meta_params, ...)
    inner_net.solve(data, ...)
    ...

    gc.collect()

Now we create the InnerNet object outside of the for-loop and reuse the instances.

inner_nets = [InnerNet(meta_params, ...) for task_id in range(num_tasks)]

for task_id in range(num_tasks):
    inner_net = inner_nets[task_id]
    inner_net.reset_parameters(...)
    inner_net.solve(data, ...)
    ...

If this solution does not resolve your issue, you can open an issue for that.

[zaccharieramzi]

I solved my issue by getting rid of torchopt.FuncOptimizer and computing grads, updates and applying them manually.

I think the problem might be that in the step function of FuncOptimizer there is a create_graph=True, but I am not too sure.