Choice equivalent for PyTorch

[LeviViana]

Related to #16897 and #18457.

[ezyang]

Let us know if you want review prior to removal of WIP

[LeviViana]

Let us know if you want review prior to removal of WIP

I think it would be nice.

[LeviViana]

I'm making some changes in the sampling_with_replacement for performance reasons. These are going to be the last improvements I have in mind prior to a review.

[LeviViana]

The CUDA extensions are not randomizing correctly. It's like the random states don't actually change that much.

import torch
x = torch.arange(10).cuda()
w = torch.arange(10).float().cuda()
# This sampling will always give either 6 or 9
torch.choice(x, w, True, 1) # gives 9
torch.choice(x, w, True, 1) # gives 6
torch.choice(x, w, True, 1) # gives 9
torch.choice(x, w, True, 1) # gives 6

If I added THCRandom_seed(state) before getting the gen_states, then it would randomize correctly, but I wouldn't validate the reproducibility tests.

I'm investigating how to fix it. I'll try to get some inspiration from the implementation of randperm.

[LeviViana]

I've made some checks available here. I won't be bringing any more changes until a review, and I'm looking forward to it! Thanks!

[soumith]

I was reviewing this PR, and my biggest comment is around implemeting choice kernels.
I think it'd be much simpler and probably much more performant if choice just used torch.multinomial + advanced indexing, instead of it's own kernels -- considering how long we took to optimize multinomial.

Also, the correctness tests look like a good start but insufficient

[LeviViana]

Thanks @soumith for your feedback. I've made some benchmarks, and I've noticed that choice is faster than multinomial + indexing.

python3 -m timeit --setup="import torch; x = torch.arange(10); w = torch.arange(10).float()" "torch.choice(x, w, False, 6)" # 11 µsec
python3 -m timeit --setup="import torch; x = torch.arange(10); w = torch.arange(10).float()" "x[torch.multinomial(w, 6, replacement=False)]" # 12 µsec

python3 -m timeit --setup="import torch; x = torch.arange(10); w = torch.arange(10).float()" "torch.choice(x, w, True, 6)" 8 µsec
python3 -m timeit --setup="import torch; x = torch.arange(10); w = torch.arange(10).float()" "x[torch.multinomial(w, 6, replacement=True)]" 11 µsec

python3 -m timeit --setup="import torch; x = torch.arange(10)" "torch.choice(x, False, 6)" 7 µsec
python3 -m timeit --setup="import torch; x = torch.arange(10); w = torch.ones(10).float()" "x[torch.multinomial(w, 6, replacement=False)]" 13 µsec

python3 -m timeit --setup="import torch; x = torch.arange(10)" "torch.choice(x, True, 6)" 4 µsec
python3 -m timeit --setup="import torch; x = torch.arange(10); w = torch.ones(10).float()" "x[torch.multinomial(w, 6, replacement=True)]" 13 µsec

I could perform more tests with bigger tensors and with CUDA later if needed. Indeed, when I first implemented choice I wasn't aware of the existence of multinomial.

I can improve the correctness tests, but first I'd like to know what would make you interested in this implementation. It seems to me that this implementation only makes sense now if I could prove performance gains in all cases.

[LeviViana]

I've done some tests with big CUDA Tensors x = torch.arange(10 ** 5), k = 10 ** 3, here are the results:

• Weighted Sampling WITHOUT replacement: 707 µsec vs 2.8 sec
• Weighted Sampling WITH replacement: 739 µsec vs 2.8 sec
• Uniform Sampling WITHOUT replacement: 233 µsec vs 3.17 sec
• Uniform Sampling WITH replacement: 12 µsec vs 2.9 sec

It looks like torch.choice is faster, but these gaps seem huge to me... I can't find any error. Here below you'll find some of the snippets to reproduce these results.

python3 -m timeit --setup="import torch; x = torch.arange(10 ** 5).cuda(); w = torch.arange(10 ** 5).float().cuda()" "torch.choice(x, w, True, 10 ** 3)" 
python3 -m timeit --setup="import torch; x = torch.arange(10 ** 5).cuda(); w = torch.arange(10 ** 5).float().cuda()" "x[torch.multinomial(w, 10 ** 3, replacement=True)]" 


python3 -m timeit --setup="import torch; x = torch.arange(10 ** 5).cuda()" "torch.choice(x, True, 10 ** 3)" 
python3 -m timeit --setup="import torch; x = torch.arange(10 ** 5).cuda(); w = torch.ones(10 ** 5).float().cuda()" "x[torch.multinomial(w, 10 ** 3, replacement=True)]"

[LeviViana]

I've been checking the choice's distribution sampling correctness. So far, everything is working properly. Here is a snippet I used to do some tests:

Change the parameters, m, n, k, replace and device.

import torch
import torch.nn.functional as F
import numpy as np

m = 20
n = 10000
k = 4
replace = False
device = 'cpu'

###################################
# Comparing Choice vs Multinomial #
###################################

multinomial_samples = []
choice_samples = []

weights = torch.rand(m, device=device)

for _ in range(n):
	multinomial_samples += torch.multinomial(
						  weights,
						  k,
						  replace
						).cpu().numpy().tolist()
	choice_samples += torch.choice(
					torch.arange(m).to(device),
					weights,
					replace,
					k
				      ).cpu().numpy().tolist()

_, multinomial_dist = np.unique(multinomial_samples, return_counts=True)
_, choice_dist = np.unique(choice_samples, return_counts=True)

multinomial_dist = torch.Tensor(multinomial_dist) / (n * k)
choice_dist = torch.Tensor(choice_dist) / (n * k)

print(F.kl_div(choice_dist.log(), multinomial_dist, reduction='sum'))

############################################
# Comparing Choice vs Correct distribution #
############################################

choice_samples = []
weights = torch.rand(m, device=device)

for _ in range(n):
	choice_samples += torch.choice(
					torch.arange(m).to(device),
					weights,
					replace,
					1
				      ).cpu().numpy().tolist()

correct_dist =  weights / weights.sum()
correct_dist = correct_dist.to('cpu')
_, choice_dist = np.unique(choice_samples, return_counts=True)
choice_dist = torch.Tensor(choice_dist) / n

print(F.kl_div(choice_dist.log(), correct_dist, reduction='sum'))

[fmassa]

I had a quick look at our current implementation of multinomial on the GPU (not alias_multinomial, which is currently private).

It seems that our GPU implementation launches n_samples kernels

pytorch/aten/src/THC/generic/THCTensorRandom.cu

Lines 252 to 274 in b6d0f6c

	for (int sample = 0; sample < n_sample; ++sample) {
	if (sample > 0) {
	// Update probabilities
	// Renorm along rows
	THCTensor_(copy)(state, normDist, origDist);
	THCTensor_(renormRows)(state, normDist);
	// Prefix sum along rows
	THCTensor_(cumsum)(state, prefixSum, normDist, 1);
	}
	// The kernel can only draw one sample before we have to
	// recalculate our distribution
	sampleMultinomialWithoutReplacement
	<<<grid, block, 0, THCState_getCurrentStream(state)>>>(
	gen->state.gen_states,
	n_sample,
	sample,
	THCudaLongTensor_data(state, self),
	numDist, numCategories,
	THCTensor_(data)(state, origDist),
	THCTensor_(data)(state, prefixSum));
	}

which might explain the slowdown of torch.multinomial compared to @LeviViana implementation of choice, which launches only one or two kernels.

I didn't think carefully why this is needed for torch.multinomial, maybe it's because torch.multinomial supports passing a matrix of probabilities, while choice only supports a vector.

Also, looking at the implementation of choice, it seems that the last step in choice is always to call index_select on the first tensor. If this is indeed the case, then it could make sense to potentially remove the index_select from the kernels, and maybe see if there is something that we could do to integrate some of the ideas from choice in the implementation of multinomial.

[umanwizard]

I would prefer if the arguments had the same names and order as in its inspiration in NumPy.

(With the exception of input, which we can leave as-is, since that's a PyTorch standard).

[LeviViana]

I'll make this change as well.

[umanwizard]

Here instead of duplicating the code you can just call through to native::choice_cpu(input, weights, replace, k)

[LeviViana]

Indeed, I'll make this change.

[umanwizard]

So, I guess we are assuming the input tensor x is 1-D, which seems reasonable since that's what's NumPy does.

But we need to actually check that and report an error if it's not the case.

[LeviViana]

I'm not assuming x is 1-D, instead I'm forcing the sampling to happen only in the first dimension. If x = torch.Tensor([[1, 2], [3, 4]]) then torch.choice(x, w, True, 3) can be torch.Tensor([[1, 2], [1, 2], [3, 4]]) for instance (i.e. it is sampling 0, 0, 1).

[umanwizard]

Got it. Anyway, we should still check that there is at least one dimension, because this will barf if you call it on a 0-dim tensor.

[umanwizard]

If somebody passes in a tensor with all zero weights, this divides by zero.

[LeviViana]

You are right, I'll fix this.

[umanwizard]

Why not allow weights to be an integral type?

[LeviViana]

I'll make that change as well. I guess some casting will be necessary in the CUDA kernels though.

[umanwizard]

why?

[LeviViana]

You are right, It doesn't. I can just check whether the weights are contiguous and in the case they aren't I can just force it to be. I'll make this change, thanks.

[fmassa]

This is looking good, thanks!

I've looked the CPU part for now, there are a few things that I think should be improved. Let me know what you think

[fmassa]

you can always make weights.contiguous() unconditionally, this will avoid a copy in the case it is already contiguous.

[fmassa]

do you need to clone it here?

[LeviViana]

Not anymore !

[fmassa]

there is a typo here. it's meant to be contiguous

[fmassa]

you have an assert in the beginning that the weights should be float, so this doesn't work for double?

[LeviViana]

You are right, I'll fix this.

[fmassa]

do you mean that you want it to be floating point type? Like double or float?

[LeviViana]

I'll remove this check.

[LeviViana]

I've done some benchmarks against numpy. I'm using only the CPU implementation for these tests. Here are the results for sampling 2k items out of 10k elements.

Edit: These benchmarks give the same results using OMP_NUM_THREADS=1.

Sampling type	torch.choice	numpy.random.choice
Weighted without replacement	550 µsec	980 µsec
Weighted with replacement	240 µsec	390 µsec
Uniform without replacement	53 µsec	191 µsec
Uniform with replacement	35 µsec	44 µsec

To reproduce:

python3 -m timeit --setup="import torch; n=10000; k=2000; x = torch.arange(n); w = torch.arange(n).float()" "torch.choice(x, k, False, w)" # 550 µsec
python3 -m timeit --setup="import numpy as np; n=10000; k=2000; x = np.arange(n); w = np.arange(n); w = w / w.sum()" "np.random.choice(x, k, False, w)" # 980 µsec

python3 -m timeit --setup="import torch; n=10000; k=2000; x = torch.arange(n); w = torch.arange(n).float()" "torch.choice(x, k, True, w)" # 240 µsec
python3 -m timeit --setup="import numpy as np; n=10000; k=2000; x = np.arange(n); w = np.arange(n); w = w / w.sum()" "np.random.choice(x, k, True, w)" # 390 µsec

python3 -m timeit --setup="import torch; n=10000; k=2000; x = torch.arange(n)" "torch.choice(x, k, False)" # 53 µsec
python3 -m timeit --setup="import numpy as np; n=10000; k=2000; x = np.arange(n)" "np.random.choice(x, k, False)" # 191 µsec

python3 -m timeit --setup="import torch; n=10000; k=2000; x = torch.arange(n)" "torch.choice(x, k, True)" # 35 µsec
python3 -m timeit --setup="import numpy as np; n=10000; k=2000; x = np.arange(n)" "np.random.choice(x, k, True)" # 44 µsec

[ptrblck]

@LeviViana I'm not sure, how python -m timeit works internally, but is it synchronizing the CUDA calls automatically or are we missing the torch.cuda.synchronize() calls in these tests:

python3 -m timeit --setup="import torch; x = torch.arange(10 ** 7).cuda(); w = torch.arange(10 ** 7).cuda().float() + 1; J, q = torch._multinomial_alias_setup(w)" "x[torch._multinomial_alias_draw(q, J, 10 ** 4)]"
python3 -m timeit --setup="import torch; x = torch.arange(10 ** 7).cuda(); w = torch.arange(10 ** 7).cuda().float() + 1" "torch.choice(x, 10 ** 4, True, w)" 

[LeviViana]

Thanks @ptrblck, you are right, torch.cuda.synchronize() is missing. I've performed some tests and the times reported by the cuda calls are underestimated. However, for the moment all the conclusions drawn so far are still valid.

[ptrblck]

Thanks for the update @LeviViana!

I've built your current branch and used this script to profile both methods:

import torch
import time

# Setup
x = torch.arange(10 ** 7).cuda()
w = torch.arange(10 ** 7).cuda().float() + 1
nb_iters = 1000

# warmup
for _ in range(50):
    J, q = torch._multinomial_alias_setup(w)
    output = x[torch._multinomial_alias_draw(q, J, 10 ** 4)]

# Profile1
torch.cuda.synchronize()
t0 = time.time()

for _ in range(nb_iters):
    J, q = torch._multinomial_alias_setup(w)
    output = x[torch._multinomial_alias_draw(q, J, 10 ** 4)]

torch.cuda.synchronize()
t1 = time.time()
print('elapsed {:.6f}s/iter'.format((t1 - t0)/nb_iters))


# warmup
for _ in range(50):
    output = torch.choice(x, 10 ** 4, True, w)

# Profile2
torch.cuda.synchronize()
t0 = time.time()

for _ in range(nb_iters):
    output = torch.choice(x, 10**4, True, w)

torch.cuda.synchronize()
t1 = time.time()
print('elapsed {:.6f}s/iter'.format((t1 - t0)/nb_iters))

On a TitanV, CUDA10.1.105 and cudnn7500 I get the following numbers:

# multinomial 
elapsed 0.002027s/iter

# choice
elapsed 0.005599s/iter

Could you run this script and check your current output?

[LeviViana]

Thanks @ptrblck for the script. These are my outputs (RTX 2080Ti, CUDA 10.1, CuDNN 7.5):

# multinomial
elapsed 0.002217s/iter

# choice
elapsed 0.004415s/iter

Indeed, the multinomial aliases are faster than choice for sampling with replacement, despite the torch._multinomial_alias_setup overhead.

[LeviViana]

It would be interesting to establish an objective to this PR. So far, what I've noticed is that the torch._multinomial_alias_draw is faster for sampling with replacement. So, I propose the following plan:

• torch.choice for sampling with replacement will just call torch._multinomial_alias_setup and torch._multinomial_alias_draw
• Keep torch.choice for sampling without replacement

This way we get the best of both implementations. If you validate this plan, I can make the changes and update the PR. What do you think @gchanan @soumith ?

[pytorchbot]

Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as Stale.
Feel free to remove the Stale label if you feel this was a mistake.
If you are unable to remove the Stale label please contact a maintainer in order to do so.
Stale pull requests will automatically be closed 30 days after being marked Stale

[kmfrick]

Any updates on this? It could be interesting to integrate this in the C++ API (I don't know if the .yaml files allow for auto-generating C++ API code too), since you have to use a slower .index() method to perform random sampling on tensors there (see comments on this issue).

[kmfrick]

Also, why is this getting compared to torch._multinomial_alias_draw when that function has been killed?

[github-actions]

Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as Stale.
Feel free to remove the Stale label if you feel this was a mistake.
If you are unable to remove the Stale label please contact a maintainer in order to do so.
If you want the bot to never mark this PR stale again, add the no-stale label.
Stale pull requests will automatically be closed after 30 days of inactivity.

[github-actions]

Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as Stale.
Feel free to remove the Stale label if you feel this was a mistake.
If you are unable to remove the Stale label please contact a maintainer in order to do so.
If you want the bot to never mark this PR stale again, add the no-stale label.
Stale pull requests will automatically be closed after 30 days of inactivity.

[facebook-github-bot]

/easycla

As part of the transition to the PyTorch Foundation, this project now requires contributions be covered under the new CLA. See #85559 for additional details.

This comment will trigger a new check of this PR. If you are already covered, you will simply see a new "EasyCLA" check that passes. If you are not covered, a bot will leave a new comment with a link to sign.

[linux-foundation-easycla]

• ❌ - login: @LeviViana / name: Levi Viana . The commit (86510e4, c7e0325, f4a7e81, 322cbdc, 1a815fc, 8ee2fbe, 0157252, 98d7385, 4364807, d3e6846, 6012f6c, 2b58385, 83156a3, 0fb56c4, 2713c95, de1395b, 5b20d07, aac9b3e, 1c60a45, c0f33eb, fb8252d, cc430ed, ba140bd, 760fd07, ec4096c, 06cebcb, b863356) is not authorized under a signed CLA. Please click here to be authorized. For further assistance with EasyCLA, please submit a support request ticket.

[ducha-aiki]

Any hope to get this re-opened?

[stupidcucumber]

Why did this issue closed?... It would be really great to have torch.choice() for clean code(. There are numerous examples where I import numpy just to implement function that does np.random.choice(). So sad to see as this beautiful PR has drown...