Reproducibility of examples

[PaulScemama]

Hi!

I am currently trying to run experiments that are reproducible, but even after controlling for data splits / shuffling I seem to not be able to get reproducible results.

For example, in the MINST classification example, consider I edit the download() function so that shuffling is fixed:

import tensorflow as tf
import tensorflow_datasets as tfds


def download(split_range, shuffle=False):
    ds = tfds.load(
        name="MNIST",
        split=f"train[{split_range}]",
        as_supervised=True,
        shuffle_files=False,
    ).map(lambda x, y: (tf.cast(x, tf.float32) / 255.0, y))
    if shuffle:
        ds = ds.shuffle(10, reshuffle_each_iteration=False)
    return ds.batch(128).prefetch(1)


train_data_loader, val_data_loader, test_data_loader = (
    download(":80%", shuffle=False),
    download("80%:90%"),
    download("90%:"),
)

If I run the rest of the cells in that example twice, I get different results. My question: what other modules do I need to interact with in order to produce deterministic results?

Thanks in advance!

[gianlucadetommaso]

Hi Paul,
I don't seem to reproduce the issue. All I did it was set shuffle to False, that is:

train_data_loader, val_data_loader, test_data_loader = (
    download(":80%", shuffle=False),
    download("80%:90%"),
    download("90%:"),
)

Then I run once and I get results. I restart the kernel. I run a second time and I seem to get the same results as for the first run.

Importantly, make sure that you re-instantiate prob_model, as there is an internal random number generator stepping forward every time it is needed.

Can you check whether this is still an issue? Can you tell me at which point you start seeing differences? Thanks!

[PaulScemama]

Hi Gianluca,

I went to lunch and now it seems to work as expected! I think maybe I didn't properly restart the kernel. I am now getting reproducible results. Sorry for the bother.

[PaulScemama]

@gianlucadetommaso I actually had to set a key to pass to the model.predictive.mean() and model.predictive.mode() to get reproducible test accuracy, brier score, and ece.

import jax
rng = jax.random.PRNGKey(0)

test_log_probs = laplace_model.predictive.log_prob(data_loader=test_data_loader)
test_inputs_loader = test_data_loader.to_inputs_loader()
test_means = laplace_model.predictive.mean(inputs_loader=test_inputs_loader, rng=rng) #RNG
test_modes = laplace_model.predictive.mode(
    inputs_loader=test_inputs_loader, means=test_means, rng=rng #RNG
)

[gianlucadetommaso]

You shouldn't need to set the rng there either! As long as you run the cells in the same sequence, you should be getting the same results. 😄

[PaulScemama]

If I run the cells in one sequence I get reproducible results. But say I run the full sequence once. Then I re-run only these lines of code

test_log_probs = laplace_model.predictive.log_prob(data_loader=test_data_loader)
test_inputs_loader = test_data_loader.to_inputs_loader()
test_means = laplace_model.predictive.mean(inputs_loader=test_inputs_loader,)
test_modes = laplace_model.predictive.mode(
    inputs_loader=test_inputs_loader, means=test_means,
)

from fortuna.metric.classification import (
    accuracy,
    expected_calibration_error,
    brier_score,
)

test_targets = test_data_loader.to_array_targets()
acc = accuracy(preds=test_modes, targets=test_targets)
brier = brier_score(probs=test_means, targets=test_targets)
ece = expected_calibration_error(
    preds=test_modes,
    probs=test_means,
    targets=test_targets,
    plot=True,
    plot_options=dict(figsize=(10, 2)),
)
print(f"Test accuracy: {acc}")
print(f"Brier score: {brier}")
print(f"ECE: {ece}")

Each time I run those lines of code over again I get different results. This matters because I want to save the posterior/model and then load it in later to evaluate.

However, if I use the rng argument:

import jax
rng = jax.random.PRNGKey(0)

test_log_probs = laplace_model.predictive.log_prob(data_loader=test_data_loader)
test_inputs_loader = test_data_loader.to_inputs_loader()
test_means = laplace_model.predictive.mean(inputs_loader=test_inputs_loader, rng=rng)
test_modes = laplace_model.predictive.mode(
    inputs_loader=test_inputs_loader, means=test_means, rng=rng
)

from fortuna.metric.classification import (
    accuracy,
    expected_calibration_error,
    brier_score,
)

test_targets = test_data_loader.to_array_targets()
acc = accuracy(preds=test_modes, targets=test_targets)
brier = brier_score(probs=test_means, targets=test_targets)
ece = expected_calibration_error(
    preds=test_modes,
    probs=test_means,
    targets=test_targets,
    plot=True,
    plot_options=dict(figsize=(10, 2)),
)
print(f"Test accuracy: {acc}")
print(f"Brier score: {brier}")
print(f"ECE: {ece}")

And I keep rerunning these cells, I get the same results.

[gianlucadetommaso]

Alright, yes this makes sense and it is expected. The predictive statistics are based on sampling from the posterior, which uses the random number generator set inside the probabilistic model. However, as you are doing, you can also pass the rng externally.

Good that you found the solution yourself!

I'll close the issue, but feel free to ask about any related questions!

[PaulScemama]

Great, thank you!