CUDA out of memory when using ot.sinkhorn2 as a loss function about pot HOT 4 CLOSED

Hello,

This probably comes form the fact that torch needs to keep the intermediate values in sinkhorn in memory to allow for a backward(). I have been planning to add a detach_iterations parameter for a hile that will allow to run the whole algorithm and plug the gradient assuming convrgence (implicit differettiation) for a while. I will get to that when I have more time.

Could you please check if sinkhorn2 explodes in memory when you give it arrays with keep_gradient=False?

from pot.

Hello, I have added an option to ot.solvethat forces the use of implicit gradients and limits the memory use:

you can compute the sinkhorn loss with

loss = ot.solve(M, a, b, reg=1, grad='implicit').value

All iterations are detached and the gradient is set at the end with no memory overhead but then it is differentiable only wrt value (not value_linear or the OT plan). Could you tell me if it solves your problem, it is merges in master brach ?

from pot.

This probably comes form the fact that torch _needs to keep the intermediate values in sinkhorn in memory to allow for a backward(). I have been planning to add a detach_iterations parameter for a hile that will allow to run the whole algorithm and plug the gradient assuming convrgence - [ ] # (implicit differettiation) for a while. I will` get to that when I have more time.**

from pot.

This new functionality is now available in ot.solve and ot.solve_sample thanks to PR #605 .

from pot.