About the computation of the EMD Loss

[qsisi]

Hello! Thanks for open-sourcing this amazing work. Here I got a question about the EMD loss.

My intention is to understand the exact way of the computation for the EMD loss. So I changed this line:

pcrnet_pytorch/pcrnet/losses/cuda/emd_torch/pkg/layer/emd_loss_layer.py

Line 12 in fad0886

return cost

to this:
return cost, match
In order to get the optimal transport matrix. Then see how to compute the cost from match.
Following:

from emd import EMDLoss
p1 = torch.rand(1, 1024, 3).cuda()
p2 = torch.rand(1, 1024, 3).cuda()
emd_loss = EMDLoss()
cost, match = emd_loss(p1, p2)

Even though I got the optimal transport matrix as match in the code, I still couldn't figure out how to compute the cost from match. To go through the cuda code is tough for me, so could you kindly help me with this?

Thank you very much for your help.