隨機數字產生

在 TensorFlow.org 上檢視 在 Google Colab 中執行 在 GitHub 上檢視原始碼 下載筆記本

TensorFlow 在 tf.random 模組中提供一組虛擬隨機數字產生器 (RNG)。本文件說明如何控制隨機數字產生器,以及這些產生器如何與其他 tensorflow 子系統互動。

TensorFlow 提供兩種方法來控制隨機數字產生程序

  1. 透過明確使用 tf.random.Generator 物件。每個此類物件都維護一個狀態 (tf.Variable 中),該狀態會在每次數字產生後變更。

  2. 透過純粹功能性的無狀態隨機函式,例如 tf.random.stateless_uniform。使用相同的引數 (包括種子) 且在相同的裝置上呼叫這些函式,一律會產生相同的結果。

設定

import tensorflow as tf

# Creates some virtual devices (cpu:0, cpu:1, etc.) for using distribution strategy
physical_devices = tf.config.list_physical_devices("CPU")
tf.config.experimental.set_virtual_device_configuration(
    physical_devices[0], [
        tf.config.experimental.VirtualDeviceConfiguration(),
        tf.config.experimental.VirtualDeviceConfiguration(),
        tf.config.experimental.VirtualDeviceConfiguration()
    ])
2024-01-17 02:22:51.386100: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2024-01-17 02:22:51.386148: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2024-01-17 02:22:51.387696: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered

tf.random.Generator 類別

tf.random.Generator 類別用於您希望每次 RNG 呼叫都產生不同結果的情況。它維護一個內部狀態 (由 tf.Variable 物件管理),該狀態會在每次產生隨機數字時更新。由於狀態由 tf.Variable 管理,因此它享有 tf.Variable 提供的所有功能,例如輕鬆檢查點、自動控制依附性和執行緒安全。

您可以手動建立類別物件,或呼叫 tf.random.get_global_generator() 以取得預設全域產生器,藉此取得 tf.random.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.5658045  -0.1012345  -0.2744976 ]], shape=(2, 3), dtype=float32)
tf.Tensor(
[[ 0.24077678  0.39891425  0.03557164]
 [-0.15206331 -0.7270625   1.8158559 ]], 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.5658045  -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(
[[-0.8503367  -0.8919918   0.688985  ]
 [-0.51400167  0.57703274 -0.5177701 ]], 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)

split 將變更呼叫它的產生器 (g 在上述範例中) 的狀態,類似於 RNG 方法,例如 normal。除了彼此獨立之外,新的產生器 (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(0.4637335, shape=(), dtype=float32)

您可以遞迴方式執行分割,在分割的產生器上呼叫 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 的引數使用時,不同的產生器物件將導致 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

請注意,此重新追蹤行為與 tf.Variable 一致

num_traces = 0
@tf.function
def foo(v):
  global num_traces
  num_traces += 1
  return v.read_value()
foo(tf.Variable(1))
foo(tf.Variable(2))
print(num_traces)
1

與分散式策略互動

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)
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():
  g = tf.random.Generator.from_seed(1)
  print(strat.run(lambda: g.normal([])))
  print(strat.run(lambda: g.normal([])))
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.
PerReplica:{
  0: tf.Tensor(-0.87930447, shape=(), dtype=float32),
  1: tf.Tensor(0.020661574, shape=(), dtype=float32)
}
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.
PerReplica:{
  0: tf.Tensor(-1.5822568, shape=(), dtype=float32),
  1: tf.Tensor(0.77539235, shape=(), dtype=float32)
}

如果產生器已植入種子 (例如由 Generator.from_seed 建立),則隨機數字由種子決定,即使不同的複本取得不同且不相關的數字。可以將複本上產生的隨機數字視為複本 ID 的雜湊,以及所有複本共有的「主要」隨機數字。因此,整個系統仍然是決定性的。

也可以在 Strategy.run 內部建立 tf.random.Generator

strat = tf.distribute.MirroredStrategy(devices=["cpu:0", "cpu:1"])
with strat.scope():
  def f():
    g = tf.random.Generator.from_seed(1)
    a = g.normal([])
    b = g.normal([])
    return tf.stack([a, b])
  print(strat.run(f))
  print(strat.run(f))
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.
PerReplica:{
  0: tf.Tensor([-0.87930447 -1.5822568 ], shape=(2,), dtype=float32),
  1: tf.Tensor([0.02066157 0.77539235], shape=(2,), dtype=float32)
}
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.
PerReplica:{
  0: tf.Tensor([-0.87930447 -1.5822568 ], shape=(2,), dtype=float32),
  1: tf.Tensor([0.02066157 0.77539235], shape=(2,), dtype=float32)
}

我們不再建議將 tf.random.Generator 作為引數傳遞至 Strategy.run,因為 Strategy.run 通常預期引數為張量,而非產生器。

儲存產生器

一般來說,對於儲存或序列化,您可以像處理 tf.Variabletf.Module (或其子類別) 一樣處理 tf.random.Generator。在 TF 中,有兩種序列化機制:檢查點SavedModel

檢查點

產生器可以使用 tf.train.Checkpoint 自由儲存和還原。還原點的隨機數字串流將與儲存點的隨機數字串流相同。

filename = "./checkpoint"
g = tf.random.Generator.from_seed(1)
cp = tf.train.Checkpoint(generator=g)
print(g.normal([]))
tf.Tensor(0.43842277, shape=(), dtype=float32)
cp.write(filename)
print("RNG stream from saving point:")
print(g.normal([]))
print(g.normal([]))
RNG stream from saving point:
tf.Tensor(1.6272374, shape=(), dtype=float32)
tf.Tensor(1.6307176, shape=(), dtype=float32)
cp.restore(filename)
print("RNG stream from restoring point:")
print(g.normal([]))
print(g.normal([]))
RNG stream from restoring point:
tf.Tensor(1.6272374, shape=(), dtype=float32)
tf.Tensor(1.6307176, shape=(), dtype=float32)

您也可以在分散式策略內儲存和還原

filename = "./checkpoint"
strat = tf.distribute.MirroredStrategy(devices=["cpu:0", "cpu:1"])
with strat.scope():
  g = tf.random.Generator.from_seed(1)
  cp = tf.train.Checkpoint(my_generator=g)
  print(strat.run(lambda: g.normal([])))
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:CPU:0', '/job:localhost/replica:0/task:0/device:CPU:1')
PerReplica:{
  0: tf.Tensor(-0.87930447, shape=(), dtype=float32),
  1: tf.Tensor(0.020661574, shape=(), dtype=float32)
}
with strat.scope():
  cp.write(filename)
  print("RNG stream from saving point:")
  print(strat.run(lambda: g.normal([])))
  print(strat.run(lambda: g.normal([])))
RNG stream from saving point:
PerReplica:{
  0: tf.Tensor(-1.5822568, shape=(), dtype=float32),
  1: tf.Tensor(0.77539235, shape=(), dtype=float32)
}
PerReplica:{
  0: tf.Tensor(-0.5039703, shape=(), dtype=float32),
  1: tf.Tensor(0.1251838, shape=(), dtype=float32)
}
with strat.scope():
  cp.restore(filename)
  print("RNG stream from restoring point:")
  print(strat.run(lambda: g.normal([])))
  print(strat.run(lambda: g.normal([])))
RNG stream from restoring point:
PerReplica:{
  0: tf.Tensor(-1.5822568, shape=(), dtype=float32),
  1: tf.Tensor(0.77539235, shape=(), dtype=float32)
}
PerReplica:{
  0: tf.Tensor(-0.5039703, shape=(), dtype=float32),
  1: tf.Tensor(0.1251838, shape=(), dtype=float32)
}

您應確保複本在儲存之前,其 RNG 呼叫歷史記錄不會發散 (例如,一個複本進行一次 RNG 呼叫,而另一個複本進行兩次 RNG 呼叫)。否則,它們的內部 RNG 狀態會發散,且 tf.train.Checkpoint (僅儲存第一個複本的狀態) 將無法正確還原所有複本。

您也可以將儲存的檢查點還原至具有不同複本數量的不同分散式策略。由於在策略中建立的 tf.random.Generator 物件只能在相同的策略中使用,因此若要還原至不同的策略,您必須在目標策略中建立新的 tf.random.Generator,並為其建立新的 tf.train.Checkpoint,如此範例所示

filename = "./checkpoint"
strat1 = tf.distribute.MirroredStrategy(devices=["cpu:0", "cpu:1"])
with strat1.scope():
  g1 = tf.random.Generator.from_seed(1)
  cp1 = tf.train.Checkpoint(my_generator=g1)
  print(strat1.run(lambda: g1.normal([])))
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:CPU:0', '/job:localhost/replica:0/task:0/device:CPU:1')
PerReplica:{
  0: tf.Tensor(-0.87930447, shape=(), dtype=float32),
  1: tf.Tensor(0.020661574, shape=(), dtype=float32)
}
with strat1.scope():
  cp1.write(filename)
  print("RNG stream from saving point:")
  print(strat1.run(lambda: g1.normal([])))
  print(strat1.run(lambda: g1.normal([])))
RNG stream from saving point:
PerReplica:{
  0: tf.Tensor(-1.5822568, shape=(), dtype=float32),
  1: tf.Tensor(0.77539235, shape=(), dtype=float32)
}
PerReplica:{
  0: tf.Tensor(-0.5039703, shape=(), dtype=float32),
  1: tf.Tensor(0.1251838, shape=(), dtype=float32)
}
strat2 = tf.distribute.MirroredStrategy(devices=["cpu:0", "cpu:1", "cpu:2"])
with strat2.scope():
  g2 = tf.random.Generator.from_seed(1)
  cp2 = tf.train.Checkpoint(my_generator=g2)
  cp2.restore(filename)
  print("RNG stream from restoring point:")
  print(strat2.run(lambda: g2.normal([])))
  print(strat2.run(lambda: g2.normal([])))
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:CPU:0', '/job:localhost/replica:0/task:0/device:CPU:1', '/job:localhost/replica:0/task:0/device:CPU:2')
RNG stream from restoring point:
PerReplica:{
  0: tf.Tensor(-1.5822568, shape=(), dtype=float32),
  1: tf.Tensor(0.77539235, shape=(), dtype=float32),
  2: tf.Tensor(0.6851049, shape=(), dtype=float32)
}
PerReplica:{
  0: tf.Tensor(-0.5039703, shape=(), dtype=float32),
  1: tf.Tensor(0.1251838, shape=(), dtype=float32),
  2: tf.Tensor(-0.58519536, shape=(), dtype=float32)
}

雖然 g1cp1g2cp2 是不同的物件,但它們透過通用檢查點檔案 filename 和物件名稱 my_generator 連結。策略之間的重疊複本 (例如上述的 cpu:0cpu:1) 將使其 RNG 串流像先前的範例中一樣正確還原。當產生器儲存在策略範圍內,並還原到任何策略範圍之外時,此保證不適用,反之亦然,因為策略外部的裝置會被視為與策略中的任何複本不同。

SavedModel

tf.random.Generator 可以儲存到 SavedModel。產生器可以在策略範圍內建立。儲存也可以在策略範圍內發生。

filename = "./saved_model"

class MyModule(tf.Module):

  def __init__(self):
    super(MyModule, self).__init__()
    self.g = tf.random.Generator.from_seed(0)

  @tf.function
  def __call__(self):
    return self.g.normal([])

  @tf.function
  def state(self):
    return self.g.state

strat = tf.distribute.MirroredStrategy(devices=["cpu:0", "cpu:1"])
with strat.scope():
  m = MyModule()
  print(strat.run(m))
  print("state:", m.state())
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:CPU:0', '/job:localhost/replica:0/task:0/device:CPU:1')
PerReplica:{
  0: tf.Tensor(-1.4154755, shape=(), dtype=float32),
  1: tf.Tensor(-0.11388441, shape=(), dtype=float32)
}
state: tf.Tensor([256   0   0], shape=(3,), dtype=int64)
with strat.scope():
  tf.saved_model.save(m, filename)
  print("RNG stream from saving point:")
  print(strat.run(m))
  print("state:", m.state())
  print(strat.run(m))
  print("state:", m.state())
INFO:tensorflow:Assets written to: ./saved_model/assets
RNG stream from saving point:
PerReplica:{
  0: tf.Tensor(-0.68758255, shape=(), dtype=float32),
  1: tf.Tensor(0.8084062, shape=(), dtype=float32)
}
state: tf.Tensor([512   0   0], shape=(3,), dtype=int64)
PerReplica:{
  0: tf.Tensor(-0.27342677, shape=(), dtype=float32),
  1: tf.Tensor(-0.53093255, shape=(), dtype=float32)
}
state: tf.Tensor([768   0   0], shape=(3,), dtype=int64)
imported = tf.saved_model.load(filename)
print("RNG stream from loading point:")
print("state:", imported.state())
print(imported())
print("state:", imported.state())
print(imported())
print("state:", imported.state())
RNG stream from loading point:
state: tf.Tensor([256   0   0], shape=(3,), dtype=int64)
tf.Tensor(-1.0359411, shape=(), dtype=float32)
state: tf.Tensor([512   0   0], shape=(3,), dtype=int64)
tf.Tensor(-0.06425078, shape=(), dtype=float32)
state: tf.Tensor([768   0   0], shape=(3,), dtype=int64)

不建議將包含 tf.random.Generator 的 SavedModel 載入分散式策略,因為複本都將產生相同的隨機數字串流 (這是因為複本 ID 在 SavedModel 的圖表中凍結)。

將分散式 tf.random.Generator (在分散式策略內建立的產生器) 載入非策略環境 (如上述範例) 也有一個注意事項。RNG 狀態將正確還原,但產生的隨機數字會與其策略中的原始產生器不同 (同樣是因為策略外部的裝置會被視為與策略中的任何複本不同)。

無狀態 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] 的整數張量。運算的結果完全由這個種子決定。

無狀態 RNG 使用的 RNG 演算法與裝置相關,這表示在不同裝置上執行的相同運算可能會產生不同的輸出。

演算法

一般

tf.random.Generator 類別和 stateless 函式都支援所有裝置上的 Philox 演算法 (寫為 "philox"tf.random.Algorithm.PHILOX)。

如果使用相同的演算法並從相同的狀態開始,則不同的裝置將產生相同的整數數字。它們也將產生「幾乎相同」的浮點數字,但不同裝置執行浮點運算的方式 (例如縮減順序) 可能會導致微小的數值差異。

XLA 裝置

在 XLA 驅動的裝置 (例如 TPU,以及啟用 XLA 時的 CPU/GPU) 上,也支援 ThreeFry 演算法 (寫為 "threefry"tf.random.Algorithm.THREEFRY)。此演算法在 TPU 上速度很快,但在 CPU/GPU 上相較於 Philox 速度較慢。

如需這些演算法的更多詳細資訊,請參閱論文 「平行隨機數字:簡單如 1、2、3」