使用tf.random.set_seed后TensorFlow结果仍无法复现的问题

我正在学习TensorFlow教程，使用版本2.11.0。教程代码本应生成可复现结果，方便验证练习是否正确。我调用tf.random.set_seed(1)来保证结果复现，但每次运行时，自定义MyDenseLayer的build方法中生成的权重矩阵W和偏置b值都不一样，最终输出结果也不一致。

我的代码 #

import tensorflow as tf
import numpy as np

class MyDenseLayer(tf.keras.layers.Layer):
    def __init__(self, n_output_nodes):
        super(MyDenseLayer, self).__init__()
        self.n_output_nodes = n_output_nodes

    def build(self, input_shape):
        d = int(input_shape[-1])
        # Define and initialize parameters: a weight matrix W and bias b
        # Note that parameter initialization is random!
        self.W = self.add_weight("weight", shape=[d, self.n_output_nodes]) # note the dimensionality
        self.b = self.add_weight("bias", shape=[1, self.n_output_nodes]) # note the dimensionality
        print("Weight matrix is {}".format(self.W))
        print("Bias vector is {}".format(self.b))

    def call(self, x):
        z = tf.add(tf.matmul(x, self.W), self.b)
        y = tf.sigmoid(z)
        return y

# Since layer parameters are initialized randomly, we will set a random seed for reproducibility
tf.random.set_seed(1)
layer = MyDenseLayer(3)
layer.build((1,2))
print(layer.call(tf.constant([[1.0,2.0]], tf.float32, shape=(1,2))))

两次运行的完整输出 #

treuss@foo:~/python/tensorflow$ python lab1_1.py 
2023-03-26 21:31:16.896212: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-03-26 21:31:18.021094: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2023-03-26 21:31:18.023462: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2023-03-26 21:31:18.023634: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2023-03-26 21:31:18.023976: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-03-26 21:31:18.024324: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2023-03-26 21:31:18.024471: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2023-03-26 21:31:18.024617: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2023-03-26 21:31:18.455771: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2023-03-26 21:31:18.455946: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2023-03-26 21:31:18.456125: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2023-03-26 21:31:18.456257: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1613] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 4656 MB memory:  -> device: 0, name: NVIDIA GeForce GTX 1060 6GB, pci bus id: 0000:01:00.0, compute capability: 6.1
Weight matrix is <tf.Variable 'weight:0' shape=(2, 3) dtype=float32, numpy=
array([[ 0.9970403 , -0.672126  , -0.00545013],
       [ 0.5411365 , -0.8570848 ,  0.5970814 ]], dtype=float32)>
Bias vector is <tf.Variable 'bias:0' shape=(1, 3) dtype=float32, numpy=array([[-0.9100063,  0.7671951, -0.9659226]], dtype=float32)>
tf.Tensor([[0.7630197  0.16532896 0.5554683 ]], shape=(1, 3), dtype=float32)
treuss@foo:~/python/tensorflow$ python lab1_1.py 
2023-03-26 21:31:21.245548: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-03-26 21:31:22.372605: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2023-03-26 21:31:22.375021: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2023-03-26 21:31:22.375175: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2023-03-26 21:31:22.375521: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-03-26 21:31:22.375840: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2023-03-26 21:31:22.375960: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2023-03-26 21:31:22.376067: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2023-03-26 21:31:22.801768: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2023-03-26 21:31:22.801935: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2023-03-26 21:31:22.802112: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2023-03-26 21:31:22.802213: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1613] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 4656 MB memory:  -> device: 0, name: NVIDIA GeForce GTX 1060 6GB, pci bus id: 0000:01:00.0, compute capability: 6.1
Weight matrix is <tf.Variable 'weight:0' shape=(2, 3) dtype=float32, numpy=
array([[ 0.72208846,  0.34211397,  0.04753423],
       [ 0.48018157,  0.9557345 , -0.19968122]], dtype=float32)>
Bias vector is <tf.Variable 'bias:0' shape=(1, 3) dtype=float32, numpy=array([[ 0.31122065, -0.81101143, -0.7763765 ]], dtype=float32)>
tf.Tensor([[0.8801311  0.80885255 0.24449256]], shape=(1, 3), dtype=float32)

问题原因与解决方法 #

原因

仅设置tf.random.set_seed()不足以完全复现Keras层的权重初始化结果：

1. Keras自身维护独立的随机种子状态，权重初始化会用到这个状态
2. GPU环境下，cuDNN的非确定性运算会引入额外随机波动

解决方法

修改代码，补充以下两处设置：

1. 使用tf.keras.utils.set_random_seed()统一设置全局、Python标准库、NumPy的随机种子
2. 开启GPU确定性运算，避免设备层面的随机差异

修改后的完整代码：

import tensorflow as tf
import numpy as np

# 统一设置全量随机种子，覆盖TensorFlow、Python、NumPy的随机源
tf.keras.utils.set_random_seed(1)
# 开启GPU确定性运算，消除设备层面的随机波动
tf.config.experimental.enable_op_determinism()

class MyDenseLayer(tf.keras.layers.Layer):
    def __init__(self, n_output_nodes):
        super(MyDenseLayer, self).__init__()
        self.n_output_nodes = n_output_nodes

    def build(self, input_shape):
        d = int(input_shape[-1])
        self.W = self.add_weight("weight", shape=[d, self.n_output_nodes])
        self.b = self.add_weight("bias", shape=[1, self.n_output_nodes])
        print("Weight matrix is {}".format(self.W))
        print("Bias vector is {}".format(self.b))

    def call(self, x):
        z = tf.add(tf.matmul(x, self.W), self.b)
        y = tf.sigmoid(z)
        return y

layer = MyDenseLayer(3)
layer.build((1,2))
print(layer.call(tf.constant([[1.0,2.0]], tf.float32, shape=(1,2))))

现在每次运行代码，都会得到完全一致的权重初始化结果和最终输出。

内容的提问来源于stack exchange，提问作者treuss