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).
Do note that this is only highly useful if you don’t 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 torch.unique currently automatically sorting the array in the cuda case.

randperm762×455 56.6 KB

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()