import torch
import time
N = int(1e7)
x = torch.LongTensor(N).random_(100000)
x = x.cuda()
torch.cuda.synchronize()
start = time.time()
unique, inv=torch.unique(x, sorted=False, return_inverse=True)
torch.cuda.synchronize()
print(time.time()-start)
which gives 2 minutes on a TitanX and 0.89 s on CPU.
it seems torch.unique runtime grows rapidely with the number of unique values on GPU?
EDIT:
i tried with rusty’s “Pytorch Unique” package and it is much faster on GPU (around 40ms, but much slower in CPU (7.19 s) ). https://github.com/rusty1s/pytorch_unique 38…