Faster triplet selector

[Lucashsmello]

Hi,

it is possible to make the triplet selector run much faster both in cpu and cuda. It run 5x faster in my setup. It consists of just converting some operations made in a for loop to a single tensor operation.

Here is some code that a used to test and compare running time:

import torch
from utils import FunctionNegativeTripletSelector, random_hard_negative
from time import time
import numpy as np

torch.manual_seed(123)
n = 256
X = torch.rand((n, 8), dtype=torch.float32)
Y = torch.randint(0, 4, size=(n,))

np.random.seed(123)
selector = ORIGINALFunctionNegativeTripletSelector(1.0, random_hard_negative)
t = time()
result1 = selector.get_triplets(X, Y)
t = time()-t
print("%s %.3fsec" % (selector.__class__.__name__, t))

np.random.seed(123)
selector = NEWFunctionNegativeTripletSelector(1.0, random_hard_negative)
t = time()
result2 = selector.get_triplets(X, Y)
t = time()-t
print("%s %.3fsec" % (selector.__class__.__name__, t))

assert((result1 == result2).all())

[adambielski]

HI @Lucashsmello - thanks for the PR. Yes, unfortunately the triplet selection is not optimized, it's been on the TODO list for quite some time but I didn't have much time to work on this repo. I'll take a look at your PR and merge if there are no problems, I know that even more optimization can be done