random_generator_test.py

# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests tf.random.Generator with distribution strategies."""

import functools
import os

from absl.testing import parameterized
from tensorflow.python.checkpoint import checkpoint as tracking_util
from tensorflow.python.compat import v2_compat
from tensorflow.python.distribute import combinations as ds_combinations
from tensorflow.python.distribute import multi_process_runner
from tensorflow.python.distribute import sharded_variable
from tensorflow.python.distribute import strategy_combinations
from tensorflow.python.distribute import values as dist_values
from tensorflow.python.distribute.coordinator import cluster_coordinator as coordinator_lib
from tensorflow.python.distribute.mirrored_strategy import MirroredStrategy
from tensorflow.python.eager import def_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import test_combinations as combinations
from tensorflow.python.framework import test_util
from tensorflow.python.module import module
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import array_ops_stack
from tensorflow.python.ops import stateful_random_ops as rng
from tensorflow.python.platform import test
from tensorflow.python.saved_model import load
from tensorflow.python.saved_model import save
from tensorflow.python.util import deprecation


def get_num_local_replicas(strat, values=None):
  strat_name = type(strat).__name__
  if "MultiWorker" in strat_name or "CollectiveAllReduceStrategy" in strat_name:
    if values is None:
      values = strat.run(lambda: constant_op.constant(0))
      values = strat.experimental_local_results(values)
    return len(values)
  else:
    return strat.num_replicas_in_sync


ps_strategies = [
    strategy_combinations.parameter_server_strategy_3worker_2ps_cpu,
    strategy_combinations.parameter_server_strategy_1worker_2ps_cpu,
    strategy_combinations.parameter_server_strategy_3worker_2ps_1gpu,
    strategy_combinations.parameter_server_strategy_1worker_2ps_1gpu,
]
all_strategies = (strategy_combinations.all_strategies +
                  strategy_combinations.multiworker_strategies +
                  ps_strategies)


def run_on_strategy(replica_fn, strat, coord):
  def distributed_fn():
    return strat.run(replica_fn)
  if coord is not None:
    results = coord.schedule(
        def_function.function(distributed_fn)).fetch()
  else:
    results = distributed_fn()
  return results


class GeneratorTest(test.TestCase, parameterized.TestCase):

  def setUp(self):
    super(GeneratorTest, self).setUp()
    v2_compat.enable_v2_behavior()

  def assertAllDifferent(self, tensors):
    """Checks that there are no duplicate elements anywhere among the tensors.

    Args:
      tensors: a list of tensors. They can have different shapes.
    """
    values = [array_ops.reshape(t, shape=[-1]) for t in tensors]
    values = array_ops.concat(values, axis=0)
    values = self.evaluate(values)
    values = values.tolist()
    self.assertAllEqual(len(values), len(set(values)))

  @test_util.run_v2_only
  def testCreateOutsideMirroredStrat(self):
    """Tests RNG/MirrorStrategy interaction #1.

    If an RNG is created outside a DS scope, all replicas will access the
    same RNG object, and accesses are serialized.
    """
    shape = [3, 4]
    dtype = dtypes.int32
    gen = rng.Generator.from_seed(1234)
    strat = MirroredStrategy(devices=["cpu:0", "cpu:1"])
    with strat.scope():

      def f():
        t1 = gen.uniform_full_int(shape=shape, dtype=dtype)
        t2 = gen.uniform_full_int(shape=shape, dtype=dtype)
        t = array_ops_stack.stack([t1, t2])
        return t

      results = strat.extended.call_for_each_replica(fn=f)
      values = results.values
      self.assertAllEqual(2, len(values))
      self.assertAllDifferent(values)

  @test_util.run_v2_only
  def testMirroredStratParaAsync(self):
    """Tests RNG/MirrorStrategy interaction #2.

    The user can create n independent RNGs outside strategy.scope(), where n
    is the number of replicas, and give one to each replica. The replicas can
    thus get different random-number streams.
    """
    shape = [3, 4]
    dtype = dtypes.int32
    gens = rng.get_global_generator().split(count=2)
    devices = ["cpu:0", "cpu:1"]
    strat = MirroredStrategy(devices=devices)
    # Use `PerReplica` to specify which `gen` is sent to which replica
    gens = dist_values.PerReplica([[g] for g in gens])
    with strat.scope():

      def f(gen):
        t1 = gen.uniform_full_int(shape=shape, dtype=dtype)
        t2 = gen.uniform_full_int(shape=shape, dtype=dtype)
        t = array_ops_stack.stack([t1, t2])
        return t

      results = strat.extended.call_for_each_replica(fn=f, args=gens)
      local_results = strat.experimental_local_results(results)
      self.assertAllEqual(2, len(local_results))
      self.assertAllDifferent(local_results)

  @ds_combinations.generate(
      combinations.combine(
          strat=all_strategies,
          mode=["eager"]))
  def testCrossReplica(self, strat):
    """Tests that RNG can be properly advanced in cross-replica context."""
    def read_values(dv):
      return [v.read_value() for v in strat.experimental_local_results(dv)]
    with strat.scope():
      g = rng.Generator.from_seed(1)
      s1 = read_values(g.state)
      g.normal([3])
      g.skip(4)
      s2 = read_values(g.state)
    self.assertNotAllEqual(s1[0], s2[0])
    self.assertEqual(len(s1), len(s2))
    for i in range(1, len(s1)):
      self.assertAllEqual(s1[0], s1[i])
      self.assertAllEqual(s2[0], s2[i])

  @ds_combinations.generate(
      combinations.combine(
          strat=all_strategies,
          mode=["eager"],
          jit_replica_fn=[False, True],
          seeded=[True, False],))
  def testDistStrat(self, strat, jit_replica_fn, seeded):
    """Tests RNG with distribution strategies."""
    strat_name = type(strat).__name__
    if "TPU" in strat_name and not jit_replica_fn:
      self.skipTest(
          "TPUStrategy requires the replica function (the function passed to "
          "strategy.run) to be decorated with tf.function")
    coord = None
    if "ParameterServer" in strat_name:
      coord = coordinator_lib.ClusterCoordinator(strat)
    creators = {
        True: functools.partial(rng.Generator.from_seed, 1234),
        False: rng.Generator.from_non_deterministic_state,
    }
    shape = [3, 4]
    dtype = dtypes.int32
    creator = creators[seeded]
    with strat.scope():
      gen = creator()
      def f():
        t1 = gen.uniform_full_int(shape=shape, dtype=dtype)
        t2 = gen.uniform_full_int(shape=shape, dtype=dtype)
        t = array_ops_stack.stack([t1, t2])
        return t
      replica_fn = def_function.function(f) if jit_replica_fn else f
      results = run_on_strategy(replica_fn, strat, coord)
      values = strat.experimental_local_results(results)
      n = get_num_local_replicas(strat, values)
      self.assertAllEqual(n, len(values))
      self.assertAllDifferent(values)

  @ds_combinations.generate(
      combinations.combine(
          strat=[
              strategy_combinations.parameter_server_strategy_fn(
                  "ParameterServer1Worker2PSCPUFixedShards",
                  num_workers=1, num_ps=2,
                  variable_partitioner=(
                      sharded_variable.FixedShardsPartitioner(2)))
          ],
          mode=["eager"]))
  def testShardedError(self, strat):
    """Tests error about sharding is raised."""
    with strat.scope():
      with self.assertRaisesRegex(
          ValueError, "state is sharded, which is not allowed"):
        rng.Generator.from_seed(1234)

  @ds_combinations.generate(
      combinations.combine(
          strat=all_strategies,
          mode=["eager"],
          jit_replica_fn=[False, True]))
  def testDistVarAsTFFunArg(self, strat, jit_replica_fn):
    """Tests that RNG with dist variables can be used as tf.function's arg."""
    strat_name = type(strat).__name__
    if "CentralStorage" in strat_name:
      self.skipTest(
          "CentralStorageStrategy wraps variable updates in merge_call which "
          "can't be called inside a tf.function that doesn't cover the entire "
          "replica function (the function passed to strategy.run).")
    if "TPU" in strat_name and not jit_replica_fn:
      self.skipTest(
          "TPUStrategy requires the replica function (the function passed to "
          "strategy.run) to be decorated with tf.function")
    coord = None
    if "ParameterServer" in strat_name:
      coord = coordinator_lib.ClusterCoordinator(strat)
    shape = [3, 4]
    dtype = dtypes.int32
    with strat.scope():
      gen = rng.Generator.from_seed(1234)
      @def_function.function
      def f(gen):  # the main focus
        t1 = gen.uniform_full_int(shape=shape, dtype=dtype)
        t2 = gen.uniform_full_int(shape=shape, dtype=dtype)
        t = array_ops_stack.stack([t1, t2])
        return t
      def g():
        return f(gen)
      replica_fn = def_function.function(g) if jit_replica_fn else g
      for _ in range(2):
        results = run_on_strategy(replica_fn, strat, coord)
        values = strat.experimental_local_results(results)
        n = get_num_local_replicas(strat, values)
        self.assertAllEqual(n, len(values))
        self.assertAllDifferent(values)

  @ds_combinations.generate(
      combinations.combine(
          strat1=strategy_combinations.all_strategies,
          strat2=strategy_combinations.all_strategies,
          jit_replica_fn=[False, True],
          mode=["eager"]) +
      combinations.combine(
          strat1=strategy_combinations.multiworker_strategies + ps_strategies,
          strat2=[None],
          jit_replica_fn=[False, True],
          mode=["eager"]))
  def testDistStratRestore(self, strat1, strat2, jit_replica_fn):
    """Tests checkpointing and restoring (to possibly different #replicas)."""
    if strat2 is None:
      strat2 = strat1
    strat1_name = type(strat1).__name__
    strat2_name = type(strat2).__name__
    if "Default" in strat1_name or "Default" in strat2_name:
      self.skipTest(
          "We don't guarantee consistency between strategy and no-strategy.")
    if ("TPU" in strat1_name or "TPU" in strat2_name) and not jit_replica_fn:
      self.skipTest(
          "TPUStrategy requires the replica function (the function passed to "
          "strategy.run) to be decorated with tf.function")
    coord1 = None
    if "ParameterServer" in strat1_name:
      coord1 = coordinator_lib.ClusterCoordinator(strat1)
    coord2 = None
    if "ParameterServer" in strat2_name:
      coord2 = coordinator_lib.ClusterCoordinator(strat2)
    fname = os.path.join(self.get_temp_dir(), "checkpoint")
    def uniform(strat, coord, g):
      def f():
        return g.uniform_full_int([3], dtype=dtypes.int32)
      replica_fn = def_function.function(f) if jit_replica_fn else f
      result = run_on_strategy(replica_fn, strat, coord)
      return strat.experimental_local_results(result)
    with strat1.scope():
      g1 = rng.Generator.from_seed(1)
    with strat2.scope():
      g2 = rng.Generator.from_seed(10)
    cp1 = tracking_util.Checkpoint(g=g1)
    cp2 = tracking_util.Checkpoint(g=g2)
    def write_restore_compare():
      cp1.write(fname)
      r1 = uniform(strat1, coord1, g1)
      cp2.restore(fname)
      r2 = uniform(strat2, coord2, g2)
      # Tests that overlapping replicas are properly restored.
      n1 = get_num_local_replicas(strat1)
      n2 = get_num_local_replicas(strat2)
      n = min(n1, n2)
      self.assertAllEqual(r1[:n], r2[:n])
    # Run multiple times so that cp1.write is called in various RNG states
    for _ in range(2):
      write_restore_compare()

  @ds_combinations.generate(
      combinations.combine(
          strat=all_strategies,
          mode=["eager"],
          is_save_in_scope=[True, False]))
  def testSavedModel(self, strat, is_save_in_scope):

    class CustomModule(module.Module):

      def __init__(self):
        super(CustomModule, self).__init__()
        self.g = rng.Generator.from_seed(0)

      @def_function.function
      def __call__(self):
        return self.g.state

      @def_function.function
      def mutate(self):
        self.g.normal([])

    with strat.scope():
      m = CustomModule()
      m.mutate()
      state_before = m()
      path = os.path.join(self.get_temp_dir(), "saved_model")
    if is_save_in_scope:
      with strat.scope():
        save.save(m, path)
    else:
      save.save(m, path)
    with strat.scope():
      m.mutate()
      state_before_2 = m()

    imported = load.load(path)
    state_after = imported()
    self.assertAllEqual(state_before, state_after)
    imported.mutate()
    state_after_2 = imported()
    self.assertAllEqual(state_before_2, state_after_2)


if __name__ == "__main__":
  with deprecation.silence():
    multi_process_runner.test_main()