Torch equivalent of numpy.random.choice?

phan_phan: <p>Hi,<br>
I am trying to extract random “slices” of tensors.<br>
Is there a torch equivalent of <strong>numpy.random.choice</strong> ?</p>
<p>And if not, is there a reason why not ?</p>
<p>Thanks</p>

richard: <p>You could generate a random number between 0 and the size of the outer dimension of your tensor, and then use that to index into your tensor.</p>
<p>We don’t have a built-in function like <code>numpy.random.choice</code>.</p>

phan_phan: <p>Thanks!<br>
Actually I wanted to draw k samples, and without replacement,<br>
so I ended up doing :<br>
<strong>perm = torch.randperm(tensor.size(0))</strong><br>
<strong>idx = perm[:k]</strong><br>
<strong>samples = tensor[idx]</strong></p>
<p>(but maybe that’s not computationally efficient)</p>

dgh328: <p>Well, the main advantage of numpy.random.choice is the possibility to pass in an array of probabilities corresponding to each element, which this solution does not cover. Even python’s random library enables passing a weight list to its choices() function.</p>

richard: <p>Oh, are you looking for <code>torch.multinomial</code>? <a href="http://pytorch.org/docs/master/torch.html?highlight=multinomial#torch.multinomial">http://pytorch.org/docs/master/torch.html?highlight=multinomial#torch.multinomial</a></p>

LeviViana: <p>There is an issue currently opened in PyTorch’s github repo about that subject: <a href="https://github.com/pytorch/pytorch/issues/16897" rel="nofollow noopener">https://github.com/pytorch/pytorch/issues/16897</a></p>

LeviViana: <p>A first version of a full-featured  <code>numpy.random.choice</code>  equivalent for PyTorch is now available <a href="https://github.com/LeviViana/torch_sampling" rel="nofollow noopener">here</a> (working on PyTorch 1.0.0). It includes CPU and CUDA implementations of:</p>
<ul>
<li>Uniform Random Sampling WITH Replacement (via  <code>torch::randint</code> )</li>
<li>Uniform Random Sampling WITHOUT Replacement (via reservoir sampling)</li>
<li>Weighted Random Sampling WITH Replacement (via inverse transform sampling)</li>
<li>Weighted Random Sampling WITHOUT Replacement (via <a href="https://utopia.duth.gr/~pefraimi/research/data/2007EncOfAlg.pdf" rel="nofollow noopener">this method</a>)</li>
</ul>

LeviViana: <p>Update: There is currently a <a href="https://github.com/pytorch/pytorch/pull/18624" rel="nofollow noopener">PR</a> waiting for review in the PyTorch’s repo.</p>

Jiatao_GU: <p>This is extremely slow and unacceptable.</p>

achark: <p>In my case: values.shape = (386363948, 2), k = 190973, the following code works fairly fast, 0.1 ~ 0.2 second.</p>
<pre><code class="lang-python">indice = random.sample(range(386363948), 190973)
indice = torch.tensor(indice)
sampled_values = values[indice]
</code></pre>
<p>Using torch.randperm, however, would cost more than 20 seconds.</p>
<pre><code class="lang-python">sampled_values = values[torch.randperm(386363948)[190973]]
</code></pre>

LeviViana: <aside class="quote no-group" data-username="achark" data-post="10" data-topic="16146">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/achark/48/12228_2.png" class="avatar"> achark:</div>
<blockquote>
<p>190973</p>
</blockquote>
</aside>
<p>Answer <a href="https://github.com/LeviViana/torch_sampling/issues/1#issuecomment-596462581" rel="noopener nofollow ugc">here</a> !</p>

anirvirnn_mishra: <p>you can do simply by applying below logic<br>
samples=torch.tensor([-11,5,9])<br>
rand_choices=samples[torch.randint(len(samples),(7,))]    #’'7 choices ‘’<br>
print(rand_choices)<br>
hope this helps!</p>

Brando_Miranda: <p>is there a reason why you prefer:</p>
<pre><code class="lang-auto">rand_choices=samples[torch.randint(len(samples),(7,))] #’'7 choices ‘’
</code></pre>
<p>vs</p>
<p><code>**perm = torch.randperm(tensor.size(0))**</code></p>
<p>?</p>
<p>Just curious</p>

phan_phan: <p>Hi !<br>
<strong>To clear things up</strong></p>
<p>If you want to do the equivalent of  <a href="https://numpy.org/doc/stable/reference/random/generated/numpy.random.choice.html" rel="noopener nofollow ugc">numpy.random.choice</a>:</p>
<pre data-code-wrap="python"><code class="lang-python">a = np.array([1, 2, 3, 4])
p = np.array([0.1, 0.1, 0.1, 0.7])
n = 2
replace = True
</code></pre>
<blockquote>
<p><strong>b = np.random.choice(a, p=p, size=n, replace=replace)</strong></p>
</blockquote>
<p>In pytorch you can use <a href="https://pytorch.org/docs/stable/generated/torch.multinomial.html#torch.multinomial" rel="noopener nofollow ugc">torch.multinomial</a> :</p>
<pre><code class="lang-auto">a = torch.tensor([1, 2, 3, 4])
p = torch.tensor([0.1, 0.1, 0.1, 0.7])
n = 2
replace = True
</code></pre>
<blockquote>
<p><strong>idx = p.multinomial(num_samples=n, replacement=replace)<br>
b = a[idx]</strong></p>
</blockquote>
<p>Careful, <code>np.random.choice</code> defaults to <code>replace=True</code><br>
But <code>torch.multinomial</code> defaults to <code>replacement=False</code></p>

111179: <p>In case the <code>*num_samples*</code> is not int type, how to deal implement the above case? Thank you!</p>

playertr: <p>If anyone is here looking for fast ways to select samples, I created a small comparison to time some of the popular random indexing solutions from the forums. For the comparison, I wrote small functions with the goal of generating indices to select 10% of a population. For a fair comparison, the indices are returned as GPU-tensors.</p>
<p>It looks like, if your population size is less than int32.MAX_VALUE, generating a random permutation on the GPU may be the fastest solution. Mileage may vary, so I’ve included my entire plotting script below so you can test it.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/8/d82fb40a006b343a20a3d92b41888cc9aeac43ce.png" data-download-href="https://discuss.pytorch.org/uploads/default/d82fb40a006b343a20a3d92b41888cc9aeac43ce" title="randperm"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/8/d82fb40a006b343a20a3d92b41888cc9aeac43ce_2_690x413.png" alt="randperm" data-base62-sha1="uQtkV28AwocQJufn2PmkVuUK250" width="690" height="413" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/8/d82fb40a006b343a20a3d92b41888cc9aeac43ce_2_690x413.png, https://discuss.pytorch.org/uploads/default/original/3X/d/8/d82fb40a006b343a20a3d92b41888cc9aeac43ce.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/d/8/d82fb40a006b343a20a3d92b41888cc9aeac43ce.png 2x" data-dominant-color="F6F5F5"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">randperm</span><span class="informations">760×455 50.2 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<pre><code class="lang-auto">import numpy as np
import torch
import random
import time
import pandas as pd
import matplotlib.pyplot as plt

def rand(pop_size, num_samples):
    """Use random.sample to generate indices."""
    return torch.Tensor(random.sample(range(pop_size), num_samples)).to('cuda')

def mult_cpu(pop_size, num_samples):
    """Use torch.Tensor.multinomial to generate indices on a CPU tensor."""
    p = torch.ones(pop_size) / pop_size
    return p.multinomial(num_samples=num_samples, replacement=False).to('cuda')

def mult_gpu(pop_size, num_samples):
    """Use torch.Tensor.multinomial to generate indices on a GPU tensor."""
    p = torch.ones(pop_size, device='cuda') / pop_size
    return p.multinomial(num_samples=num_samples, replacement=False)

def perm_cpu(pop_size, num_samples):
    """Use torch.randperm to generate indices on a CPU tensor."""
    return torch.randperm(pop_size)[:num_samples].to('cuda')

def perm_gpu(pop_size, num_samples):
    """Use torch.randperm to generate indices on a GPU tensor."""
    return torch.randperm(pop_size, device='cuda')[:num_samples]

def perm_gpu_f32(pop_size, num_samples):
    """Use torch.randperm to generate indices on a 32-bit GPU tensor."""
    return torch.randperm(pop_size, dtype=torch.int32, device='cuda')[:num_samples]

def sort_rand_cpu(pop_size, num_samples):
    """Generate a random torch.Tensor (CPU) and sort it to generate indices."""
    return torch.argsort(torch.rand(pop_size))[:num_samples]

def sort_rand_gpu(pop_size, num_samples):
    """Generate a random torch.Tensor (GPU) and sort it to generate indices."""
    return torch.argsort(torch.rand(pop_size, device='cuda'))[:num_samples]

idx_fns = [rand, 
    mult_cpu, 
    mult_gpu, 
    perm_cpu, 
    perm_gpu,
    perm_gpu_f32,
    sort_rand_cpu,
    sort_rand_gpu]

pop_size = np.logspace(3, 7, 50, dtype=int)

d = []
for n_p in pop_size:
    print(f"Testing functions with {n_p} points.")
    for fn in idx_fns:
        tic = time.time()
        num_samples = int(0.1 * n_p)
        samples = fn(n_p, num_samples)
        toc = time.time()
        assert type(samples) == torch.Tensor

        d.append({
            'Population': n_p,
            'Samples': num_samples,
            'Function':fn.__name__,
            'Time': toc - tic
        })

df = pd.DataFrame(d)

fig, ax = plt.subplots()

for fn in idx_fns:
    idxs = df['Function'] == fn.__name__
    ax.plot(df['Population'][idxs], df['Time'][idxs], label=fn.__name__)

ax.set_title("Runtime of Index-Generating Methods for Randomly Selecting 10% of Population")
ax.set_xlabel('Population Size')
ax.set_ylabel('Time')
ax.set_yscale('log')
ax.legend(bbox_to_anchor=(1, 0.7))
plt.savefig('figs/randperm.png', bbox_inches='tight')
plt.show()
</code></pre>

ptrblck: <p>Thanks for sharing the code!<br>
Since GPU operations are executed asynchronously, you would have to synchronize the code manually before starting and stopping the timer via <code>torch.cuda.synchronize()</code> to get the real execution time.<br>
Otherwise you might be profiling the kernel launch times and blocking operations would accumulate the execution time of already running kernels.<br>
<code>torch.utils.benchmark</code> provides a utility to run such comparisons and will add warmup iterations and the needed synchronizations for you.</p>

playertr: <p>Ooh, thanks! I’ll have a look and see if I can update with proper benchmarks when I have a minute.</p>

playertr: <p>Here are the results  with proper benchmarks! It looks like generating random permutations on the GPU is still the way to go, if you want to generate indices for random selection. However, the GPU methods do not scale quite as well as it seemed before.</p>
<p>I couldn’t find a good way to access the benchmark results, so I settled for <code>timeit(N).raw_times[0]</code>, which seems to give the median time spent.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/f/7f4a8fc7a8d75d1e9e65dce746717d17a2ab68e2.png" data-download-href="https://discuss.pytorch.org/uploads/default/7f4a8fc7a8d75d1e9e65dce746717d17a2ab68e2" title="randperm"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/f/7f4a8fc7a8d75d1e9e65dce746717d17a2ab68e2_2_690x413.png" alt="randperm" data-base62-sha1="ia4nX1YRn7zb9gCl7VDc4KwAfiG" width="690" height="413" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/f/7f4a8fc7a8d75d1e9e65dce746717d17a2ab68e2_2_690x413.png, https://discuss.pytorch.org/uploads/default/original/3X/7/f/7f4a8fc7a8d75d1e9e65dce746717d17a2ab68e2.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/7/f/7f4a8fc7a8d75d1e9e65dce746717d17a2ab68e2.png 2x" data-dominant-color="F7F5F5"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">randperm</span><span class="informations">760×455 47.9 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<pre><code class="lang-auto">import numpy as np
import torch
import random
import pandas as pd
import matplotlib.pyplot as plt
import torch.utils.benchmark as benchmark

def rand(pop_size, num_samples):
    """Use random.sample to generate indices."""
    return torch.Tensor(random.sample(range(pop_size), num_samples)).to('cuda')

def mult_cpu(pop_size, num_samples):
    """Use torch.Tensor.multinomial to generate indices on a CPU tensor."""
    p = torch.ones(pop_size) / pop_size
    return p.multinomial(num_samples=num_samples, replacement=False).to('cuda')

def mult_gpu(pop_size, num_samples):
    """Use torch.Tensor.multinomial to generate indices on a GPU tensor."""
    p = torch.ones(pop_size, device='cuda') / pop_size
    return p.multinomial(num_samples=num_samples, replacement=False)

def perm_cpu(pop_size, num_samples):
    """Use torch.randperm to generate indices on a CPU tensor."""
    return torch.randperm(pop_size)[:num_samples].to('cuda')

def perm_gpu(pop_size, num_samples):
    """Use torch.randperm to generate indices on a GPU tensor."""
    return torch.randperm(pop_size, device='cuda')[:num_samples]

def perm_gpu_f32(pop_size, num_samples):
    """Use torch.randperm to generate indices on a 32-bit GPU tensor."""
    return torch.randperm(pop_size, dtype=torch.int32, device='cuda')[:num_samples]

def sort_rand_cpu(pop_size, num_samples):
    """Generate a random torch.Tensor (CPU) and sort it to generate indices."""
    return torch.argsort(torch.rand(pop_size))[:num_samples]

def sort_rand_gpu(pop_size, num_samples):
    """Generate a random torch.Tensor (GPU) and sort it to generate indices."""
    return torch.argsort(torch.rand(pop_size, device='cuda'))[:num_samples]

idx_fns = [rand, 
    mult_cpu, 
    mult_gpu, 
    perm_cpu, 
    perm_gpu,
    perm_gpu_f32,
    sort_rand_cpu,
    sort_rand_gpu]

pop_size = np.logspace(3, 7, 10, dtype=int)

d = []
for n_p in pop_size:
    print(f"Testing functions with {n_p} points.")
    for fn in idx_fns:
        num_samples = int(0.1 * n_p)

        t0 = benchmark.Timer(
            stmt=f"{fn.__name__}(n_p, num_samples)",
            setup=f"from __main__ import {fn.__name__}",
            globals={'n_p': n_p, 'num_samples': num_samples}
        )

        d.append({
            'Population': n_p,
            'Samples': num_samples,
            'Function':fn.__name__,
            'Time': t0.timeit(11).raw_times[0]
        })

df = pd.DataFrame(d)

fig, ax = plt.subplots()

for fn in idx_fns:
    idxs = df['Function'] == fn.__name__
    ax.plot(df['Population'][idxs], df['Time'][idxs], label=fn.__name__)

ax.set_title("Runtime of Index-Generating Methods for Randomly Selecting 10% of Population")
ax.set_xlabel('Population Size')
ax.set_ylabel('Time')
ax.set_yscale('log')
ax.legend(bbox_to_anchor=(1, 0.7))
plt.savefig('figs/randperm.png', bbox_inches='tight')
plt.show()
</code></pre>

Filip_Cornell: <p>I had the same problem and came up with an additional way to implement my own, and it seems to work fairly well actually. I ran it with the code above (I’ll post my exact code below).</p>
<p>Do note that this is only highly useful if you <strong>don’t</strong> care about having random shuffles, but rather just random slices. If you want random shuffles, it has the same speed as randperm, more or less. This is due to <code>torch.unique</code> currently automatically sorting the array in the cuda case.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/f/af40d58bd18eb05ef6864cb056a93e6309ad9841.png" data-download-href="https://discuss.pytorch.org/uploads/default/af40d58bd18eb05ef6864cb056a93e6309ad9841" title="randperm"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/f/af40d58bd18eb05ef6864cb056a93e6309ad9841_2_690x412.png" alt="randperm" data-base62-sha1="p0msLaZbBlVO2J3yL7SL6z24VYR" width="690" height="412" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/f/af40d58bd18eb05ef6864cb056a93e6309ad9841_2_690x412.png, https://discuss.pytorch.org/uploads/default/original/3X/a/f/af40d58bd18eb05ef6864cb056a93e6309ad9841.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/a/f/af40d58bd18eb05ef6864cb056a93e6309ad9841.png 2x" data-dominant-color="F5F5F4"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">randperm</span><span class="informations">762×455 56.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<pre data-code-wrap="python3"><code class="lang-python3">import numpy as np
import torch
import random
import pandas as pd
import matplotlib.pyplot as plt
import torch.utils.benchmark as benchmark


def rand(pop_size, num_samples):
    """Use random.sample to generate indices."""
    return torch.Tensor(random.sample(range(pop_size), num_samples)).to("cuda")


def mult_cpu(pop_size, num_samples):
    """Use torch.Tensor.multinomial to generate indices on a CPU tensor."""
    p = torch.ones(pop_size) / pop_size
    return p.multinomial(num_samples=num_samples, replacement=False).to("cuda")


def mult_gpu(pop_size, num_samples):
    """Use torch.Tensor.multinomial to generate indices on a GPU tensor."""
    p = torch.ones(pop_size, device="cuda") / pop_size
    return p.multinomial(num_samples=num_samples, replacement=False)


def perm_cpu(pop_size, num_samples):
    """Use torch.randperm to generate indices on a CPU tensor."""
    return torch.randperm(pop_size)[:num_samples].to("cuda")


def perm_gpu(pop_size, num_samples):
    """Use torch.randperm to generate indices on a GPU tensor."""
    return torch.randperm(pop_size, device="cuda")[:num_samples]


def perm_gpu_f32(pop_size, num_samples):
    """Use torch.randperm to generate indices on a 32-bit GPU tensor."""
    return torch.randperm(pop_size, dtype=torch.int32, device="cuda")[:num_samples]


def sort_rand_cpu(pop_size, num_samples):
    """Generate a random torch.Tensor (CPU) and sort it to generate indices."""
    return torch.argsort(torch.rand(pop_size))[:num_samples]


def sort_rand_gpu(pop_size, num_samples):
    """Generate a random torch.Tensor (GPU) and sort it to generate indices."""
    return torch.argsort(torch.rand(pop_size, device="cuda"))[:num_samples]


def torchrand(pop_size, num_samples, device: str):
    vec = torch.unique(
        (torch.rand(num_samples, device=device) * pop_size).floor().long()
    )
    # Eliminate all duplicate entries. Might slow down the procedure but totally worth it.
    while vec.shape[0] != num_samples:
        vec = torch.unique(
            torch.cat(
                [
                    vec,
                    (torch.rand(num_samples - vec.shape[0], device=device) * pop_size)
                    .floor()
                    .long(),
                ]
            )
        )

    return vec.view(-1)


def torchrand_cpu(pop_size, num_samples, shuffle=False):
    return (
        torchrand(pop_size, num_samples, device="cpu")
        if shuffle is False
        else torch.randperm(pop_size)
    )  # [:num_samples]


def torchrand_gpu(pop_size, num_samples, shuffle=False):
    return (
        torchrand(pop_size, num_samples, device="cuda")
        if shuffle is False
        else torch.randperm(pop_size)
    )  # [:num_samples]


idx_fns = [
    rand,
    mult_cpu,
    mult_gpu,
    perm_cpu,
    perm_gpu,
    perm_gpu_f32,
    sort_rand_cpu,
    sort_rand_gpu,
    torchrand_cpu,
    torchrand_gpu,
]

pop_size = np.logspace(3, 7, 10, dtype=int)
d = []
for n_p in pop_size:
    print(f"Testing functions with {n_p} points.")
    for fn in idx_fns:
        num_samples = int(0.1 * n_p)

        t0 = benchmark.Timer(
            stmt=f"{fn.__name__}(n_p, num_samples)",
            setup=f"from __main__ import {fn.__name__}",
            globals={"n_p": n_p, "num_samples": num_samples},
        )

        d.append(
            {
                "Population": n_p,
                "Samples": num_samples,
                "Function": fn.__name__,
                "Time": t0.timeit(11).raw_times[0],
            }
        )

df = pd.DataFrame(d)

fig, ax = plt.subplots()

for fn in idx_fns:
    idxs = df["Function"] == fn.__name__
    ax.plot(df["Population"][idxs], df["Time"][idxs], label=fn.__name__)

ax.set_title(
    "Runtime of Index-Generating Methods for Randomly Selecting 10% of Population"
)
ax.set_xlabel("Population Size")
ax.set_ylabel("Time")
ax.set_yscale("log")
ax.legend(bbox_to_anchor=(1, 0.7))
plt.savefig("figs/randperm.png", bbox_inches="tight")
plt.show()


</code></pre>