How to get GPU accleration for solving Gromov Wasserstain problem?

[rucnyz]

I am using gromov_wasserstein function in my project and try to get acceleration on my GPU server. However this method is numpy-based and I have to rewrite it to meet my need. Finally I find it hard to modify the hardcoded emd_c method.

I search Google but don't get help about "Solve Earth Movers Distance problem in pytorch". I know this is not a bug or enhancement and I'm sorry for that, but I really don't know where to ask after these attempts. Do you have any suggestions?

[rflamary]

Hello, I converted this to discussion that is a better place.

GW solver use the exact OT solver in C++ so it will not be accelerated in GPU. As stated in the documentation the data will be copied back to CPU to solve the problem and then converted back to GPU. There is to the best of my knowledge no exact OT solve run GPU but using directly pytorch array on emd2 or gromov_wasserstein2 function will return a differentiable loss on CPU or GPU.

If you need GPU acceleration you can use entropic GW that provide smooth losses and will run fully on GPU.

[rucnyz]

Thank you for your reply!
However it seems that sinkhorn performs really bad in my work, and I have to use optim.cg for better performance:(

[rflamary]

sinkhorn is slow because convergence criterion is low by default (or regularization is low). Most people use sinkhorn with just a few iterations (which does not converge) which is enough in most applications. For entropic GW it is hard to find a good balance indeed.

[rucnyz]

yes, I tried grid search for regularization term and used large enough numItermax, but still get nan for too low and constant for too high. Maybe entropic regularization is a good solution for ot but not for GW problem.