无需更换的样品

Question

如何在 TensorFlow 中进行无替换采样？就像numpy.random.choice(n, size=k, replace=False) 用于一些非常大的整数n（例如100k-100M）和更小的k（例如100-10k）。 此外，我希望它高效并在 GPU 上运行，因此像 this 和 tf.py_func 这样的其他解决方案对我来说并不是一个真正的选择。任何会使用tf.range(n) 左右的东西也不是一个选择，因为n 可能非常大。

Answer 1

这是一种方式：

n = ...
sample_size = ...
idx = tf.random_shuffle(tf.range(n))[:sample_size]

---

编辑：

我已经在下面发布了答案，但随后阅读了您帖子的最后一行。如果你绝对不能产生大小为 O(n) 的张量，我认为没有什么好的方法（numpy.random.choice 和 replace=False 也被实现为排列的一部分）。您可以求助于tf.while_loop，直到您拥有唯一索引：

n = ...
sample_size = ...
idx = tf.zeros(sample_size, dtype=tf.int64)
idx = tf.while_loop(
    lambda i: tf.size(idx) == tf.size(tf.unique(idx)),
    lambda i: tf.random_uniform(sample_size, maxval=n, dtype=int64))

---

编辑 2：

关于前一种方法的平均迭代次数。如果我们称 n 为可能值的数量，而 k 为所需向量的长度（其中 k ≤ n )，一次迭代成功的概率为：

p = product((n - (i - 1) / n) 对于 i 在 1 .. k)

由于每个迭代都可以视为Bernoulli trial，因此首次成功的平均试验次数为 1 / p (proof here)。这是一个计算 Python 中某些 k 和 n 值的平均试验次数的函数：

def avg_iter(k, n):
    if k > n or n <= 0 or k < 0:
        raise ValueError()
    avg_it = 1.0
    for p in (float(n) / (n - i) for i in range(k)):
        avg_it *= p
    return avg_it

以下是一些结果：

+-------+------+----------+
|   n   |  k   | Avg iter |
+-------+------+----------+
|    10 |    5 | 3.3      |
|   100 |   10 | 1.6      |
|  1000 |   10 | 1.1      |
|  1000 |  100 | 167.8    |
| 10000 |   10 | 1.0      |
| 10000 |  100 | 1.6      |
| 10000 | 1000 | 2.9e+22  |
+-------+------+----------+

你可以看到它根据参数变化很大。

虽然我能想到的唯一算法是 O(k²)，但可以在固定数量的步骤中构造一个向量。在纯 Python 中是这样的：

import random

def sample_wo_replacement(n, k):
    sample = [0] * k
    for i in range(k):
        sample[i] = random.randint(0, n - 1 - len(sample))
    for i, v in reversed(list(enumerate(sample))):
        for p in reversed(sample[:i]):
            if v >= p:
                v += 1
        sample[i] = v
    return sample

random.seed(100)
print(sample_wo_replacement(10, 5))
# [2, 8, 9, 7, 1]
print(sample_wo_replacement(10, 10))
# [6, 5, 8, 4, 0, 9, 1, 2, 7, 3]

这是在 TensorFlow 中执行此操作的一种可能方式（不确定是否最好）：

import tensorflow as tf

def sample_wo_replacement_tf(n, k):
    # First loop
    sample = tf.constant([], dtype=tf.int64)
    i = 0
    sample, _ = tf.while_loop(
        lambda sample, i: i < k,
        # This is ugly but I did not want to define more functions
        lambda sample, i: (tf.concat([sample,
                                      tf.random_uniform([1], maxval=tf.cast(n - tf.shape(sample)[0], tf.int64), dtype=tf.int64)],
                                     axis=0),
                           i + 1),
        [sample, i], shape_invariants=[tf.TensorShape((None,)), tf.TensorShape(())])
    # Second loop
    def inner_loop(sample, i):
        sample_size = tf.shape(sample)[0]
        v = sample[i]
        j = i - 1
        v, _ = tf.while_loop(
            lambda v, j: j >= 0,
            lambda v, j: (tf.cond(v >= sample[j], lambda: v + 1, lambda: v), j - 1),
            [v, j])
        return (tf.where(tf.equal(tf.range(sample_size), i), tf.tile([v], (sample_size,)), sample), i - 1)
    i = tf.shape(sample)[0] - 1
    sample, _ = tf.while_loop(lambda sample, i: i >= 0, inner_loop, [sample, i])
    return sample

还有一个例子：

with tf.Graph().as_default(), tf.Session() as sess:
    tf.set_random_seed(100)
    sample = sample_wo_replacement_tf(10, 5)
    for i in range(10):
        print(sess.run(sample))
# [3 0 6 8 4]
# [5 4 8 9 3]
# [1 4 0 6 8]
# [8 9 5 6 7]
# [7 5 0 2 4]
# [8 4 5 3 7]
# [0 5 7 4 3]
# [2 0 3 8 6]
# [3 4 8 5 1]
# [5 7 0 2 9]

不过，这对 tf.while_loops 来说是相当深入的，众所周知，它在 TensorFlow 中并不是特别快，所以如果没有某种基准测试，我不知道使用这种方法到底能达到多快。

---

编辑 4：

最后一种可能的方法。您可以在大小为 c 的“块”中划分可能值的范围（0 到 n），并从每个块中选择随机数量的数字，然后将所有内容随机播放。您使用的内存量受 c 限制，并且您不需要嵌套循环。如果 n 可以被 c 整除，那么你应该得到一个完美的随机分布，否则最后一个“短”块中的值将获得一些额外的概率（这可能可以忽略不计视情况而定）。这是一个 NumPy 实现。考虑不同的极端情况和陷阱有点长，但是如果 c ≥ k 并且 n mod c = 0 几个部分得到简化。

import numpy as np

def sample_chunked(n, k, chunk=None):
    chunk = chunk or n
    last_chunk = chunk
    parts = n // chunk
    # Distribute k among chunks
    max_p = min(float(chunk) / k, 1.0)
    max_p_last = max_p
    if n % chunk != 0:
        parts += 1
        last_chunk = n % chunk
        max_p_last = min(float(last_chunk) / k, 1.0)
    p = np.full(parts, 2)
    # Iterate until a valid distribution is found
    while not np.isclose(np.sum(p), 1) or np.any(p > max_p) or p[-1] > max_p_last:
        p = np.random.uniform(size=parts)
        p /= np.sum(p)
    dist = (k * p).astype(np.int64)
    sample_size = np.sum(dist)
    # Account for rounding errors
    while sample_size < k:
        i = np.random.randint(len(dist))
        while (dist[i] >= chunk) or (i == parts - 1 and dist[i] >= last_chunk):
            i = np.random.randint(len(dist))
        dist[i] += 1
        sample_size += 1
    while sample_size > k:
        i = np.random.randint(len(dist))
        while dist[i] == 0:
            i = np.random.randint(len(dist))
        dist[i] -= 1
        sample_size -= 1
    assert sample_size == k
    # Generate sample parts
    sample_parts = []
    for i, v in enumerate(np.nditer(dist)):
        if v <= 0:
            continue
        c = chunk if i < parts - 1 else last_chunk
        base = chunk * i
        sample_parts.append(base + np.random.choice(c, v, replace=False))
    sample = np.concatenate(sample_parts, axis=0)
    np.random.shuffle(sample)
    return sample

np.random.seed(100)
print(sample_chunked(15, 5, 4))
# [ 8  9 12 13  3]

sample_chunked(100000000, 100000, 100000) 的快速基准测试在我的计算机中大约需要 3.1 秒，而我无法使用相同的参数运行之前的算法（上面的sample_wo_replacement 函数）来完成。应该可以在 TensorFlow 中实现它，也许可以使用 tf.TensorArray，尽管它需要付出巨大的努力才能完全正确。

Answer 2

在这里使用 gumbel-max 技巧：https://github.com/tensorflow/tensorflow/issues/9260

z = -tf.log(-tf.log(tf.random_uniform(tf.shape(logits),0,1))) 
_, indices = tf.nn.top_k(logits + z,K)

指数是你想要的。这个打勾好简单~！

Answer 3

以下在 GPU 上运行相当快，并且在使用 n~100M 和 k~10k（使用 NVIDIA GeForce GTX 1080 Ti）时我没有遇到内存问题：

def random_choice_without_replacement(n, k):
  """equivalent to 'numpy.random.choice(n, size=k, replace=False)'"""
  return tf.math.top_k(tf.random.uniform(shape=[n]), k, sorted=False).indices