添加自注意力层后CNN模型训练结果无法复现的问题求助

我在简单CNN网络中加入自注意力层NLBlock后，重复训练会得到不同结果，模型丧失了确定性；但移除NLBlock后，训练结果又能稳定复现。我已经设置了全局随机种子，启用了TensorFlow的确定性操作，甚至尝试过用tf.matmul()和softmax实现自定义注意力块，但问题依然存在。

相关代码如下：

os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
seed = 42
os.environ['TF_DETERMINISTIC_OPS'] = '1'
tf.config.experimental.enable_op_determinism()

os.environ['PYTHONHASHSEED'] = str(seed)
os.environ['TF_CUDNN_DETERMINISTIC'] = '1' 

random.seed(seed)
np.random.seed(seed)
tf.random.set_seed(seed)

tf.keras.backend.set_floatx('float64')



class NLBlock(Layer):
    def __init__(self, num_channels, **kwargs):
        super(NLBlock, self).__init__(**kwargs)
        self.num_channels = num_channels
        self.theta = Conv1D(filters=num_channels , kernel_size=1, strides=1, padding="same")
        self.phi = Conv1D(filters=num_channels , kernel_size=1, strides=1, padding="same")
        self.g = Conv1D(filters=num_channels, kernel_size=1, strides=1, padding="same")
        self.attention_layer = Attention()  # Keras Attention Layer

    def call(self, inputs):
        # Transform feature map
        query = self.theta(inputs)  # Query (Q)
        key = self.phi(inputs)      # Key (K)
        value = self.g(inputs)      # Value (V)

        # Apply Attention Layer
        attention_output = self.attention_layer([query, key, value])

        # Residual connection
        return inputs + attention_output
        
# Define Model with NL Attention
def build_model():
    optimizer = Adam(learning_rate=0.002, beta_1=0.89, beta_2=0.995)

    inputs = tf.keras.Input(shape=(num_time_steps, 1))  # Input layer
    
    # Conv Block 1
    x = Conv1D(filters=64, kernel_size=3, activation='relu', padding='same')(inputs)
    x = BatchNormalization()(x)
    x = NLBlock(num_channels=64)(x)  # NL Attention
    x = MaxPooling1D(pool_size=2)(x)

    # Conv Block 2
    x = Conv1D(filters=32, kernel_size=3, activation='relu', padding='same')(x)
    x = BatchNormalization()(x)
    x = NLBlock(num_channels=32)(x)  # NL Attention
    x = MaxPooling1D(pool_size=2)(x)
    x = Dropout(0.1)(x)

    # Conv Block 3
    x = Conv1D(filters=16, kernel_size=3, activation='relu', padding='same')(x)
    x = BatchNormalization()(x)
    x = NLBlock(num_channels=16)(x)  # NL Attention
    x = MaxPooling1D(pool_size=2)(x)
    x = Dropout(0.35)(x)

    # Fully Connected Layers
    x = Flatten()(x)
    x = Dense(40, activation='relu')(x)
    x = Dropout(0.35)(x)
    outputs = Dense(20)(x)  # Output layer for regression
    
    # Compile Model
    model = tf.keras.Model(inputs, outputs)
    model.compile(optimizer=optimizer, loss=root_mean_squared_error, metrics=['mean_absolute_error'])
    
    return model

model = build_model()

reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.6, patience=25, min_lr=1e-6)
early_stopping = EarlyStopping(monitor='val_loss', patience=30, restore_best_weights=True)
# Step 5: Train the Model
history = model.fit(X_train_scaled, y_train, validation_data=(X_val_scaled, y_val),
                epochs=180, batch_size=90, callbacks=[reduce_lr, early_stopping], verbose=0)

test_loss, test_mae = model.evaluate(X_test_scaled, y_test, batch_size=len(X_test_scaled), verbose=0)

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