生成随机数

在 TensorFlow.org 上查看 在 Google Colab 中运行 在 GitHub 上查看源代码 下载笔记本

TensorFlow 在 tf.random 模块中提供了一组伪随机数生成器 (RNG)。本文介绍如何控制随机数生成器,以及这些生成器如何与其他 Tensorflow 子系统交互。

TensorFlow 提供了两种方法来控制随机数生成过程:

  1. 通过明确使用 tf.random.Generator 对象。每个此类对象都会在 tf.Variable 中维护一个状态,该状态在每次生成随机数后都会发生改变。

  2. 通过使用纯函数式无状态随机函数,如 tf.random.stateless_uniform。在同一设备上调用具有相同参数(包括种子)的这些函数会产生相同的结果。

警告:目前尚未弃用 TF 1.x 中的旧版 RNG(如 tf.random.uniformtf.random.normal),但强烈建议不要使用。

警告:不保证随机数在不同 TensorFlow 版本间一致,请参阅:版本兼容性

设置

import tensorflow as tf

# Creates 2 virtual devices cpu:0 and cpu:1 for using distribution strategy
physical_devices = tf.config.experimental.list_physical_devices("CPU")
tf.config.experimental.set_virtual_device_configuration(
    physical_devices[0], [
        tf.config.experimental.VirtualDeviceConfiguration(),
        tf.config.experimental.VirtualDeviceConfiguration()
    ])

tf.random.Generator

当您希望每次调用 RNG 都产生不同的结果时,可以使用 tf.random.Generator 类。它会维护一个内部状态(由 tf.Variable 对象管理),该状态在每次生成随机数时都会更新。由于该状态由 tf.Variable 管理,因此,它可以利用 tf.Variable 提供的所有功能,如简单的检查点、自动控制依赖项和线程安全性。

通过手动创建 tf.random.Generator类的一个对象,您可以获得该生成器,或者通过调用 tf.random.get_global_generator(),您可以获得默认全局生成器:

g1 = tf.random.Generator.from_seed(1)
print(g1.normal(shape=[2, 3]))
g2 = tf.random.get_global_generator()
print(g2.normal(shape=[2, 3]))
tf.Tensor(
[[ 0.43842277 -0.53439844 -0.07710262]
 [ 1.5658046  -0.1012345  -0.2744976 ]], shape=(2, 3), dtype=float32)
tf.Tensor(
[[ 0.40719825  0.12406099  0.4640385 ]
 [ 0.7706627  -0.27837673  0.14459121]], shape=(2, 3), dtype=float32)

有多种方法可以创建生成器对象。最简单的方法是使用 Generator.from_seed(代码如上),从种子创建生成器。种子可以是任何非负整数,from_seed 还有一个可选参数 alg,这是该生成器将使用的 RNG 算法。

g1 = tf.random.Generator.from_seed(1, alg='philox')
print(g1.normal(shape=[2, 3]))
tf.Tensor(
[[ 0.43842277 -0.53439844 -0.07710262]
 [ 1.5658046  -0.1012345  -0.2744976 ]], shape=(2, 3), dtype=float32)

有关详细信息,请参阅后文中的算法部分。

创建生成器的另一种方法是使用 Generator.from_non_deterministic_state。以这种方式创建的生成器首先会处于非确定状态,具体取决于时间和操作系统等因素。

g = tf.random.Generator.from_non_deterministic_state()
print(g.normal(shape=[2, 3]))
tf.Tensor(
[[ 1.0782914   0.07794157 -0.31712356]
 [-0.42656502  0.01992213 -0.5324843 ]], shape=(2, 3), dtype=float32)

还有其他方法可以创建生成器,比如说通过显式状态创建,本指南不作赘述。

当使用 tf.random.get_global_generator 来获取全局生成器时,需要注意设备放置。第一次调用 tf.random.get_global_generator 时就会创建全局生成器(从非确定状态),并将其放置在该调用的作用域内的默认设备上。举个例子,如果第一次调用 tf.random.get_global_generator 的位置在 tf.device("gpu") 作用域内,则会将全局生成器放置在 GPU 上,如果稍后要从 CPU 使用全局生成器,则会将其从 GPU 复制到 CPU。

另外还有一个函数 tf.random.set_global_generator,可用于将全局生成器替换为另一个生成器对象。使用该函数前要三思,因为 tf.function 可能已获得旧全局生成器(作为弱引用),所以,替换它可能导致它被回收,从而中断 tf.function。有一种更好的方法可以重置全局生成器,即使用一个“重置”函数(如 Generator.reset_from_seed),这样就不会创建新的生成器对象。

g = tf.random.Generator.from_seed(1)
print(g.normal([]))
print(g.normal([]))
g.reset_from_seed(1)
print(g.normal([]))
tf.Tensor(0.43842277, shape=(), dtype=float32)
tf.Tensor(1.6272374, shape=(), dtype=float32)
tf.Tensor(0.43842277, shape=(), dtype=float32)

创建独立的随机数流

许多应用都需要多个独立的随机数流,所谓独立,就是指不能相互重叠,也不能有统计学上可检测到的相关性。通过使用 Generator.split 创建多个一定相互独立的生成器即可实现此目的(即生成独立流)。

g = tf.random.Generator.from_seed(1)
print(g.normal([]))
new_gs = g.split(3)
for new_g in new_gs:
  print(new_g.normal([]))
print(g.normal([]))
tf.Tensor(0.43842277, shape=(), dtype=float32)
tf.Tensor(2.536413, shape=(), dtype=float32)
tf.Tensor(0.33186463, shape=(), dtype=float32)
tf.Tensor(-0.07144657, shape=(), dtype=float32)
tf.Tensor(-0.79253083, shape=(), dtype=float32)

normal 之类的 RNG 方法类似,split 会改变调用它的生成器的状态(上例中为 g)。除相互之间保持独立外,新生成器 (new_gs) 还一定独立于旧生成器 (g)。

当您想要确保使用的生成器位于与其他计算相同的设备上,从而避免跨设备复制的开销时,生成新生成器也很有用。例如:

with tf.device("cpu"):  # change "cpu" to the device you want
  g = tf.random.get_global_generator().split(1)[0]  
  print(g.normal([]))  # use of g won't cause cross-device copy, unlike the global generator
tf.Tensor(-1.0297492, shape=(), dtype=float32)

注:在理论上,此处可以使用 from_seed(而不是 split)之类的构造函数获取新生成器,但这样做无法保证新生成器与全局生成器相互独立。同时也有使用同一种子或导致产生重叠随机数流的种子意外创建两个生成器的风险。

您可以在拆分的生成器上调用 split,从而以递归方式执行拆分。递归深度没有限制(除非发生整数溢出)。

tf.function 交互

tf.function 一起使用时,tf.random.Generator 遵循与 tf.Variable 相同的原则。这包括三个方面:

tf.function 的外部创建生成器

tf.function 可以使用在其外部创建的生成器。

g = tf.random.Generator.from_seed(1)
@tf.function
def foo():
  return g.normal([])
print(foo())
tf.Tensor(0.43842277, shape=(), dtype=float32)

调用该函数时,用户需要确保生成器对象仍处于活动状态(没有被回收)。

tf.function 的内部创建生成器

只有 tf.function 第一次运行时,才可以在其内部创建生成器。

g = None
@tf.function
def foo():
  global g
  if g is None:
    g = tf.random.Generator.from_seed(1)
  return g.normal([])
print(foo())
print(foo())
tf.Tensor(0.43842277, shape=(), dtype=float32)
tf.Tensor(1.6272374, shape=(), dtype=float32)

将生成器作为参数传递给 tf.function

当用作 tf.function 的参数时,具有相同状态大小(状态大小由 RNG 算法确定)的不同生成器对象不会导致回溯 tf.function,而具有不同状态大小的不同生成器对象则会导致回溯。

num_traces = 0
@tf.function
def foo(g):
  global num_traces
  num_traces += 1
  return g.normal([])
foo(tf.random.Generator.from_seed(1))
foo(tf.random.Generator.from_seed(2))
print(num_traces)
2

与分布策略交互

Generator 与分布策略有三种交互方式。

在分布策略的外部创建生成器

如果是在策略作用域的外部创建的生成器,则会序列化访问此生成器的所有副本,因此,每一个副本都会得到不同的随机数。

g = tf.random.Generator.from_seed(1)
strat = tf.distribute.MirroredStrategy(devices=["cpu:0", "cpu:1"])
with strat.scope():
  def f():
    print(g.normal([]))
  results = strat.run(f)
WARNING:tensorflow:There are non-GPU devices in `tf.distribute.Strategy`, not using nccl allreduce.
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:CPU:0', '/job:localhost/replica:0/task:0/device:CPU:1')
WARNING:tensorflow:Using MirroredStrategy eagerly has significant overhead currently. We will be working on improving this in the future, but for now please wrap `call_for_each_replica` or `experimental_run` or `run` inside a tf.function to get the best performance.
tf.Tensor(0.43842274, shape=(), dtype=float32)
tf.Tensor(1.6272374, shape=(), dtype=float32)

请注意,这种使用方法可能产生性能问题,因为生成器的设备与副本不同。

在分布策略的内部创建生成器

不允许在策略作用域内部创建生成器,因为这会导致在如何复制生成器方面出现歧义。比方说,是应该复制生成器,从而让每一个副本都获得相同的随机数,还是应该“拆分”,从而让每一个副本获得不同的随机数。

strat = tf.distribute.MirroredStrategy(devices=["cpu:0", "cpu:1"])
with strat.scope():
  try:
    tf.random.Generator.from_seed(1)
  except ValueError as e:
    print("ValueError:", e)
WARNING:tensorflow:There are non-GPU devices in `tf.distribute.Strategy`, not using nccl allreduce.
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:CPU:0', '/job:localhost/replica:0/task:0/device:CPU:1')

请注意,Strategy.run 会在策略作用域内隐式运行参数函数:

strat = tf.distribute.MirroredStrategy(devices=["cpu:0", "cpu:1"])
def f():
  tf.random.Generator.from_seed(1)
try:
  strat.run(f)
except ValueError as e:
  print("ValueError:", e)
WARNING:tensorflow:There are non-GPU devices in `tf.distribute.Strategy`, not using nccl allreduce.
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:CPU:0', '/job:localhost/replica:0/task:0/device:CPU:1')
WARNING:tensorflow:Using MirroredStrategy eagerly has significant overhead currently. We will be working on improving this in the future, but for now please wrap `call_for_each_replica` or `experimental_run` or `run` inside a tf.function to get the best performance.

将生成器作为参数传递给 Strategy.run

如果您希望每个副本都使用自己的生成器,则需要通过复制或拆分创建 nn 表示副本数量)个生成器,然后将其作为参数传递给 Strategy.run

strat = tf.distribute.MirroredStrategy(devices=["cpu:0", "cpu:1"])
gs = tf.random.get_global_generator().split(2)
# to_args is a workaround for the absence of APIs to create arguments for 
# run. It will be replaced when such APIs are available.
def to_args(gs):  
  with strat.scope():
    def f():
      return [gs[tf.distribute.get_replica_context().replica_id_in_sync_group]]
    return strat.run(f)
args = to_args(gs)
def f(g):
  print(g.normal([]))
results = strat.run(f, args=args)
WARNING:tensorflow:There are non-GPU devices in `tf.distribute.Strategy`, not using nccl allreduce.
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:CPU:0', '/job:localhost/replica:0/task:0/device:CPU:1')
WARNING:tensorflow:Using MirroredStrategy eagerly has significant overhead currently. We will be working on improving this in the future, but for now please wrap `call_for_each_replica` or `experimental_run` or `run` inside a tf.function to get the best performance.
WARNING:tensorflow:Using MirroredStrategy eagerly has significant overhead currently. We will be working on improving this in the future, but for now please wrap `call_for_each_replica` or `experimental_run` or `run` inside a tf.function to get the best performance.
tf.Tensor(1.1344033, shape=(), dtype=float32)
tf.Tensor(0.4836465, shape=(), dtype=float32)

无状态 RNG

无状态 RNG 的使用方法非常简单。因为它们是纯函数,不涉及状态或副作用。

print(tf.random.stateless_normal(shape=[2, 3], seed=[1, 2]))
print(tf.random.stateless_normal(shape=[2, 3], seed=[1, 2]))
tf.Tensor(
[[ 0.5441101   0.20738031  0.07356433]
 [ 0.04643455 -1.30159    -0.95385665]], shape=(2, 3), dtype=float32)
tf.Tensor(
[[ 0.5441101   0.20738031  0.07356433]
 [ 0.04643455 -1.30159    -0.95385665]], shape=(2, 3), dtype=float32)

每个无状态 RNG 都需要一个 seed 参数,该参数必须是形状为 [2] 的整数张量。该运算的结果完全由种子确定。

算法

基本信息

tf.random.Generator 类和 stateless 函数在所有设备上都支持 Philox 算法(写作 "philox"tf.random.Algorithm.PHILOX)。

如果使用相同的算法且从相同的状态开始,则不同的设备会生成相同的整数。它们还可以生成“几乎相同”的浮点数,虽然由于设备执行浮点计算的方式不同(如降阶),数值可能存在微小的差异。

XLA 设备

在 XLA 驱动的设备(如 TPU 以及启用 XLA 时的 CPU/GPU)上,还支持 ThreeFry 算法(写作 "threefry"tf.random.Algorithm.THREEFRY)。与 Philox 算法相比,该算法在 TPU 上执行速度较快,而在 CPU/GPU 上执行速度较慢。

有关这些算法的更多详细信息,请参阅论文“Parallel Random Numbers: As Easy as 1, 2, 3”