Random Choice with Pytorch?

torch has no equivalent implementation of np.random.choice(), see the discussion here. The alternative is indexing with a shuffled index or random integers.

To do it with replacement:

  1. Generate n random indices
  2. Index your original tensor with these indices
pictures[torch.randint(len(pictures), (10,))]  

To do it without replacement:

  1. Shuffle the index
  2. Take the n first elements
indices = torch.randperm(len(pictures))[:10]

pictures[indices]

Read more about torch.randint and torch.randperm. Second code snippet is inspired by this post in PyTorch Forums.


torch.multinomial provides equivalent behaviour to numpy's random.choice (including sampling with/without replacement):

# Uniform weights for random draw
unif = torch.ones(pictures.shape[0])

idx = unif.multinomial(10, replacement=True)
samples = pictures[idx]
samples.shape
>>> torch.Size([10, 28, 28, 3])

For this size of tensor:

N, D = 386363948, 2
k = 190973
values = torch.randn(N, D)

The following code works fairly fast. It takes around 0.2s:

indices = torch.tensor(random.sample(range(N), k))
indices = torch.tensor(indices)
sampled_values = values[indices]

Using torch.randperm, however, would take more than 20s:

sampled_values = values[torch.randperm(N)[:k]]