如何创建映射随机数的tf.data.Dataset并保持随机状态一致？

这个问题我之前也碰到过，本质是tf.data的map操作对随机数生成的处理方式导致的：每次map被触发时，里面的随机数生成操作都会重新执行，所以d2和d3各自的map会生成独立的随机值，自然就出现了断言失败的情况。下面给你两种简单有效的解决方案：

方案一：预先生成随机数序列，与原数据集绑定 #

把随机数生成移到Dataset外部，预先生成好所有需要的随机值，再和原数据集zip在一起，这样整个流程中每个元素对应的随机数是固定的，不管后续怎么map都复用同一个值：

import tensorflow as tf
import numpy as np

# 预先生成与原数据集长度匹配的随机数序列
random_vals = tf.random_uniform((100,))
# 原数据集
d1 = tf.data.Dataset.from_tensor_slices(tf.zeros(100))
# 将随机数序列与原数据集绑定
d_with_rand = tf.data.Dataset.zip((d1, tf.data.Dataset.from_tensor_slices(random_vals)))

# 现在map时使用预先生成的随机数，避免重复生成
d2 = d_with_rand.map(lambda x, r: x + r)
d3 = d2.map(lambda x: x)
d4 = tf.data.Dataset.zip((d2, d3))

it = d4.make_one_shot_iterator()
fetch = it.get_next()

with tf.Session() as sess:
    try:
        while True:
            X = sess.run(fetch)
            if X[0] != X[1]:
                raise RuntimeError("Not same.")
    except tf.errors.OutOfRangeError:
        print("所有元素处理完成，断言全部通过！")

方案二：在同一个map操作中生成并复用随机数 #

如果不想预先生成所有随机数，可以在单次map操作中生成随机数，同时返回两个相同的结果，这样后续直接使用这个成对的结果，避免多次map触发重复生成：

import tensorflow as tf
import numpy as np

def generate_and_reuse_rand(x):
    # 仅生成一次随机数
    z = tf.random_uniform(())
    x_plus_z = x + z
    # 返回两个完全相同的结果
    return x_plus_z, x_plus_z

d1 = tf.data.Dataset.from_tensor_slices(tf.zeros(100))
# 直接在map中生成成对结果，无需后续再zip两个数据集
d4 = d1.map(generate_and_reuse_rand)

it = d4.make_one_shot_iterator()
fetch = it.get_next()

with tf.Session() as sess:
    try:
        while True:
            X = sess.run(fetch)
            if X[0] != X[1]:
                raise RuntimeError("Not same.")
    except tf.errors.OutOfRangeError:
        print("所有元素处理完成，断言全部通过！")

原理说明 #

tf.random_uniform这类随机数生成操作，每次被执行时都会基于当前的随机状态生成新值。在你的原始代码中，d2和d3是两个独立的Dataset，它们各自的map操作会独立触发随机数生成，所以同一个元素对应的随机值会不一样。而上面两种方案都是确保每个元素的随机数只生成一次，后续所有操作都复用这个值，自然就能保证X[0]和X[1]相等了。

内容的提问来源于stack exchange，提问作者autonomous.beta