model.evaluate() loss metric calculation in Bayesian Neural Net

Question

I am currently trying to set up a Bayesian Neural Net. Hereby, I came across a strange thing. When compiling my model with the metric keras.metrics.RootMeanSquaredError() and then letting model.evaluate() run, the function gives me the RMSE. However, when I manually calculate the RMSE, the result is way different and indeed better (and more reasonable comparing to an ANN).

Does anyone of you have an idea, why the results differ and how is model.evaluate() calculating the RMSE in a probabilistic Neural Net?

Thanks for you help!!

def create_probabilistic_bnn_model(input_dim, hidden_units, kl_weight):
    inputs = layers.Input(shape=input_dim)
    features = layers.BatchNormalization()(inputs)
    features = features

    # Create hidden layers with weight uncertainty using the DenseVariational layer.
    for units in hidden_units:
        features = tfp.layers.DenseVariational(
            units=units,
            make_prior_fn=prior,
            make_posterior_fn=posterior,
            kl_weight=kl_weight,
            activation="sigmoid",
        )(features)
        #features = tf.keras.layers.Dropout(0.5)(features, training=True)
        # if last_output is not None:
        #     features += last_output  # residual layer
        # last_output = features

    # Create a probabilistic output (Normal distribution), and use the `Dense` layer
    # to produce the parameters of the distribution.
    # We set units=2 to learn both the mean and the variance of the Normal distribution.
    distribution_params = layers.Dense(units=2)(features)

    outputs = tfp.layers.IndependentNormal(1)(distribution_params)

    model = keras.Model(inputs=inputs, outputs=outputs)
    return model


def negative_loglikelihood(targets, estimated_distribution):
    return -estimated_distribution.log_prob(targets)


for n in [8]:  # [8, 16, 24]
    for layer in [[n, n]]:  # [[n, n], [n, n, n], [n, n, n, n]]
        bnn_model_probabilistic = create_probabilistic_bnn_model(
            data_size, layer, 1 / len(train_dataloader.dataset)
        )

        bnn_model_probabilistic.compile(
            optimizer=keras.optimizers.Adam(learning_rate=0.01),
            loss=negative_loglikelihood,
            metrics=[keras.metrics.RootMeanSquaredError()],
        )

        callback = tf.keras.callbacks.EarlyStopping(
            monitor="loss", patience=PATIENCE, verbose=1
        )
        # Define the learning rate scheduler callback
        lr_callback = LearningRateScheduler(lr_scheduler, verbose=0)

        print("Start training the model...")
        history = bnn_model_probabilistic.fit(
            x=feature_train, y=label_train, epochs=EPOCHS, callbacks=            
        [callback, lr_callback],
        validation_data=[feature_val, label_val]
        )
        print("Model training finished.")
        
        print()

        bnn_model_probabilistic_history = history

print(bnn_model_probabilistic.summary())
print(bnn_model_probabilistic.metrics_names)

# list all data in history
print(history.history.keys())

_, rmse_train_bnn = bnn_model_probabilistic.evaluate(x=feature_train,     
    y=label_train, verbose=0)
_, rmse_test_bnn = bnn_model_probabilistic.evaluate(x=feature_test, 
    y=label_test, verbose=0) 


#RMSE Test
BNN_residual_test = label_test - BNN_result_dataframe_test['mean']
sq_BNN_residual_test = BNN_residual_test**2
BNN_MSE_test = np.mean(sq_BNN_residual_test)
rmse_test_bnn = np.sqrt(BNN_MSE_test)