因为在你的情况下你想要一个均匀的分布你可以加速它

idx = torch.sort(torch.randint(
    0, 128 - 15, (64, 16), device=device
), axis=1).values + torch.arange(0, 16, device=device).reshape(1, -1)
sampled_emb_user = emb_user[torch.arange(len(emb_user)).unsqueeze(-1), idx]

而不是

idx = torch.multinomial(torch.ones(64, 128, device=device), 16)
sampled_emb_user = emb_user[torch.arange(len(emb_user)).unsqueeze(-1), idx]

我的机器上的运行时是427 µs和784 µs与device='cpu';以及469 µs与device='cuda'的组合。
它是如何工作的？
排序随机数给出了具有替换的多项式分布的指数，即递增，加上arange项使其严格递增，从而消除了替换。
用一个小箱子说明
第一个
下面是一个可能更快的解决方案，但不是来自完全相同的发行版：

idx = torch.cumsum(torch.diff(
torch.sort(torch.randint(
    0, 128 - 16, (64, 17), device=device
), axis=1).values
, axis=1) + 1, axis=1) - 1
sampled_emb_user = emb_user[torch.arange(len(emb_user)).unsqueeze(-1), idx]

还有一种方法，我期望更接近确切的方法，不是非常严格的分析。

# 1-rand() to include 1 and exclude zero.

d = torch.cumsum(1 - torch.rand(64, 17, device=device
), axis=1)

# this produces a sorted tensor with values in the range [0:128-16]

d = (((128 - 15) * d[:, :-1]) / d[:, -1:]).to(torch.long)
idx = d + torch.arange(0, 16, device=device).reshape(1, -1)

但最终它往往比使用sort的方法慢。