distributed wrapper does not compute the local rank loss

[mensaochun]

Now distributed loss wrapper computes a loss for the global batch, it is global loss and every rank has the same loss, it may be not right. Is it better to compute a loss using local batch vs global batch, and every lank has its own loss，what would you think about?

[KevinMusgrave]

The gradients per rank are different. See #10 (comment)

That said, I think it can be made more efficient. I have an issue open for this: #204. Actually I spent some time trying that idea a while ago, and I think the gradients ended up being different from the current implementation and I'm not sure why.

[mensaochun]

I modify the distributed implementation of nxtent loss，it works and costs less gpu memory(batch size 32->64)

from pytorch_metric_learning.utils import loss_and_miner_utils as lmu
from pytorch_metric_learning.losses.ntxent_loss import NTXentLoss
from pytorch_metric_learning.utils.distributed import all_gather_embeddings_labels

def convert_to_pairs(indices_tuple, labels, ref_labels):
    if indices_tuple is None:
        return lmu.get_all_pairs_indices(labels, ref_labels)
    elif len(indices_tuple) == 4:
        return indices_tuple
    else:
        a, p, n = indices_tuple
        return a, p, a, n

class DistributedNTXentLoss(NTXentLoss):
    def __init__(self, **kwargs):
        super(DistributedNTXentLoss, self).__init__(**kwargs)

    def compute_loss(self, embeddings, labels, indices_tuple):
        # get all rank data
        all_embeddings, all_labels = all_gather_embeddings_labels(embeddings, labels)
        # add ref_labels
        indices_tuple = convert_to_pairs(indices_tuple, labels, ref_labels=all_labels)
        if all(len(x) <= 1 for x in indices_tuple):
            return self.zero_losses()
        # modify distance computation
        # mat = self.distance(embeddings)
        mat = self.distance(embeddings, all_embeddings)
        return self.loss_method(mat, None, indices_tuple)

[KevinMusgrave]

Thanks for the code. I'll test this for the general distributed loss wrapper

[KevinMusgrave]

But it won't be for a couple of weeks because I'm busy trying to meet a deadline

[KevinMusgrave]

I added an efficient argument, see the docs: https://kevinmusgrave.github.io/pytorch-metric-learning/distributed/

It's available in v1.0.0:

pip install pytorch-metric-learning==1.0.0