Usage of numpy.random.Generator to deal with Data-Loader Parallelism about framework-reproducibility HOT 7 CLOSED

train_dataset = (
    train_dataset
    .cache()
    .map(read_tfrecord, num_parallel_calls=AUTOTUNE)
    .shuffle(512)
    .map(
        lambda o, e: (resize_and_rescale(o), resize_and_rescale(e)),
        num_parallel_calls=AUTOTUNE
    )
    .batch(cfg.training.batch_size , drop_remainder=True)
    .map(
        lambda o, e: augment(o, e, training=True), num_parallel_calls=AUTOTUNE
    )
    .prefetch(AUTOTUNE)
)
val_dataset = (
    val_dataset
    .cache()
    .map(read_tfrecord, num_parallel_calls=AUTOTUNE)
    .map(
        lambda o, e: (resize_and_rescale(o), resize_and_rescale(e)),
        num_parallel_calls=AUTOTUNE
    )
    .batch(cfg.training.batch_size, drop_remainder=False)
    .map(lambda o, e: augment(o, e, training=False), num_parallel_calls=AUTOTUNE)
    .prefetch(AUTOTUNE)
)

import tensorflow as tf
import numpy as np

np.set_printoptions(precision=2, floatmode='fixed')

def augment_random(x):
  random = np.float32(np.random.uniform())
  # The addition here represents a computationally-expensive set of operations
  return np.float32(x) + random

def random_param():
  return np.float32(np.random.uniform())

def augment(x, y):
  # The addition here represents a computationally-expensive set of operations
  return tf.add(tf.cast(x, tf.float32), y)

def nondeterministic_pipeline():
  np.random.seed(123)
  dataset = tf.data.Dataset.range(1, 6)
  dataset = dataset.map(
      # From tf.data.Dataset::map documentation: Note that use of
      # tf.numpy_function or tf.py_function in general precludes the possibility
      # of executing user-defined transformations in parallel (because of the
      # Python GIL).
      lambda x: tf.numpy_function(augment_random, inp=[x], Tout=tf.float32),
      num_parallel_calls=5)
  return np.array(list(dataset.as_numpy_iterator()))

def deterministic_pipeline():
  np.random.seed(123)
  dataset = tf.data.Dataset.range(1, 6)
  dataset = dataset.map(
      # From tf.data.Dataset::map documentation: Note that use of
      # tf.numpy_function or tf.py_function in general precludes the possibility
      # of executing user-defined transformations in parallel (because of the
      # Python GIL).
      lambda x: (x, tf.numpy_function(random_param, inp=[], Tout=tf.float32)),
      num_parallel_calls=1)
  dataset = dataset.map(lambda x, y: augment(x, y), num_parallel_calls=5)
  return np.array(list(dataset.as_numpy_iterator()))

result1 = nondeterministic_pipeline()
result2 = nondeterministic_pipeline()
result3 = deterministic_pipeline()
result4 = deterministic_pipeline()

print("\nGenerate random parameters in parallel:")
print("run 1: ",result1)
print("run 2: ",result2)

print("\nGenerate random parameters in series:")
print("run 1: ",result3)
print("run 2: ",result4)

# Generate random parameters in parallel:
# run 1:  [1.55 2.72 3.29 4.70 5.23]
# run 2:  [1.72 2.23 3.55 4.29 5.70]

# Generate random parameters in series:
# run 1:  [1.70 2.29 3.23 4.55 5.72]
# run 2:  [1.70 2.29 3.23 4.55 5.72]