TripletMarginLoss Performance

[AlexSchuy]

Hi,

I've been trying to use TripletMarginLoss for a project, but I find that it reduces the speed of my training by 5x. For example:

from pytorch_metric_learning.losses import TripletMarginLoss
import torch

embeddings = torch.rand(4000, 8)
labels = torch.randint(0, 30, (4000,))

criterion = TripletMarginLoss(triplets_per_anchor=1)

criterion(embeddings, labels)

The final line takes ~2 seconds to compute using a Tesla V100 and Intel Xeon Gold 6148. Digging into the code, I find that the issue is the function pytorch_metric_learning.utils.loss_and_miner_utils.get_random_triplet_indices. Continuing the above example, get_random_triplet_indices(labels, t_per_anchor=1) also takes ~2 seconds.

I profiled the above function and found that 90% of the time is spent in these 3 repeated calls to np.random.choice and np.where (in bold):

# sample triplets, with a weighted distribution if weights is specified.
def get_random_triplet_indices(labels, ref_labels=None, t_per_anchor=None, weights=None):
    [...]
    for i, label in enumerate(labels):
        all_pos_pair_mask = ref_labels == label
        if ref_labels_is_labels:
            all_pos_pair_mask &= indices != i
        all_pos_pair_idx = np.where(all_pos_pair_mask)[0]
        curr_label_count = len(all_pos_pair_idx)
        if curr_label_count == 0:
            continue
        k = curr_label_count if t_per_anchor is None else t_per_anchor

        if weights is not None and not np.any(np.isnan(weights[i])):
            n_idx += c_f.NUMPY_RANDOM.choice(batch_size, k, p=weights[i]).tolist()
        else:
            possible_n_idx = list(np.where(ref_labels != label)[0])
            n_idx += c_f.NUMPY_RANDOM.choice(possible_n_idx, k).tolist()

        a_idx.extend([i] * k)
        curr_p_idx = c_f.safe_random_choice(all_pos_pair_idx, k)
        p_idx.extend(curr_p_idx.tolist())
    [...]

It seems that the issue is that these calls are within a python for loop, so there is no vectorization. Perhaps this code can be optimized? I see two solutions:

1. Write a C implementation of this function.
2. Vectorize the numpy operations and eliminate the python for loop.

I would be happy to attempt to submit a PR on this, if there is interest.

[KevinMusgrave]

Yes that function is inefficient, and having a vectorized version would be great. A PR would be very much appreciated! 😄 (Related issue: #157)

Were you able to try triplets_per_anchor="all" and compare performance? It looks like you're using a very large batch (size 4000) so I guess there isn't enough memory for that.

[KevinMusgrave]

I think the ideal version would use torch operations instead of numpy, and no for loops as you mentioned.

[KevinMusgrave]

One approach could be to call get_all_triplets_indices, and then randomly select from its output.

[AlexSchuy]

Okay, I'll give it a shot :). And yes, I have memory issues with triplets_per_anchor='all'. And actually, I use a very small batch size, I just have a slightly different use-case than what I believe is typically expected; I'm using the triplet margin loss within an image on individual points, in an attempt to embed the points such that they can then be clustered into object instances within the image.

On the proposed solution, it would not work for my use-case due to the aforementioned memory constraints.

[KevinMusgrave]

Now that I think about it more, it might make sense to remove triplets_per_anchor (it's only used in 2 classes), and instead have a similar init argument for BaseTupleMiner. Then for any miner, the user could specify "all" pairs/triplets or a number. I'll create a separate issue for this idea.

[KevinMusgrave]

The usage of distance.compute_mat is another factor to consider, especially in your case where the number of embeddings is large. See #181. The fix for this is a simple if-statement, but determining the threshold for when to use distance.pairwise_distance instead of distance.compute_mat is the difficult part.

[AlexSchuy]

So, the plan to move the triplets_per_anchor setting to the miners doesn't change the need to optimize get_random_triplet_indices, right? It sounds like you just plan on calling this function from BaseTupleMiner instead of from TripletMarginLoss

[KevinMusgrave]

Probably. If it's moved to BaseTupleMiner, then the order of operations would be:

1. Mine triplets
2. Randomly select N (or maybe N / anchor) from the mined triplets

So if your version of get_random_triplet_indices can operate on the output of get_all_triplets_indices, then it should be transferable.

[AlexSchuy]

I suppose the issue is to what degree you want the miners to be responsible for selecting triplets, vs. the losses. By moving triplets_per_anchor to the miners, I think the approach you favor is to let the miners handle selection, with the losses just having a default so that they can be used without a miner; essentially, the losses have a default miner that mines all possible tuples.

Someone with a large-scale use like mine should look for a miner that can reduce the computational complexity. Having looked at the miners a bit more, it seems they fall into two broad categories:

1. Miners that evaluate some criterion on all possible tuples, and return a subset.
2. Miners that evaluate some criterion on a subset of possible tuples, and return a subset.

The miners of type [1] tend to use get_all_triplets_indices, while the miners of type [2] tend to use get_random_triplet_indices, and so are suitable for large-scale use. There are also the 'subset batch miners', which I suppose can fall into either category, but are intended to be fed into another miner before being used by a loss function.

On a similar vein, I think that there are two primary purposes for the (non-subset batch) miners:

1. To improve training convergence and ultimately the quality of the embedding by changing the loss function, e.g., by focusing on hard samples.
2. To reduce the computational cost of the loss function by reducing the number of tuples.

To me, the focus in this package has mostly been on the first, but the second is also important for some use-cases, and is often mentioned, even in the literature that is cited in the documentation for this package (e.g., https://arxiv.org/pdf/1706.07567.pdf or https://arxiv.org/pdf/1704.01285.pdf). Some of the existing miners, such as DistanceWeightedMiner should be suitable for the second case, provided that get_random_triplet_indices can be optimized.

If I understand your proposal correctly, you suggest adding to all miners the option to randomly select a subset of the output of that miner. This could be useful for the first case, however it would not help in the second; in the second case, the random selection needs to happen before the mining.

In the context of chaining miners together, perhaps it makes sense instead to simply have a stochastic subset miner, that randomly samples a subset of the tuples? Then, a user could choose to either put it before or after another miner, depending upon their use-case.

[AlexSchuy]

I created a new issue #194 that suggests adding a section to the docs on computational optimization, to avoid more people asking questions like mine.

[KevinMusgrave]

I suppose the issue is to what degree you want the miners to be responsible for selecting triplets, vs. the losses. By moving triplets_per_anchor to the miners, I think the approach you favor is to let the miners handle selection, with the losses just having a default so that they can be used without a miner; essentially, the losses have a default miner that mines all possible tuples.

Yes

If I understand your proposal correctly, you suggest adding to all miners the option to randomly select a subset of the output of that miner. This could be useful for the first case, however it would not help in the second; in the second case, the random selection needs to happen before the mining.

In the context of chaining miners together, perhaps it makes sense instead to simply have a stochastic subset miner, that randomly samples a subset of the tuples? Then, a user could choose to either put it before or after another miner, depending upon their use-case.

This probably makes the most sense. So instead of adding a tuples_per_anchor argument to BaseMiner, you pass in the mined tuples to a RandomPairMiner or RandomTripletMiner. And if these miners are used by themselves, then they randomly select from all possible tuples. This chaining behavior could be implemented in a similar way to HDCMiner (which I believe is the only chainable miner in this library).

[AlexSchuy]

Okay, so I'll try to work on this and submit a PR this weekend.