xref: /external/tensorflow/tensorflow/python/data/kernel_tests/test_base.py

# Copyright 2018 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.
# ==============================================================================
"""Test utilities for tf.data functionality."""
import os
import random
import re

from tensorflow.python.data.experimental.ops import lookup_ops as data_lookup_ops
from tensorflow.python.data.experimental.ops import random_access
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.util import nest
from tensorflow.python.data.util import structure
from tensorflow.python.eager import context
from tensorflow.python.framework import combinations
from tensorflow.python.framework import config
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_dataset_ops
from tensorflow.python.ops import gen_experimental_dataset_ops
from tensorflow.python.ops import lookup_ops
from tensorflow.python.ops import tensor_array_ops
from tensorflow.python.ops.ragged import ragged_tensor
from tensorflow.python.platform import test


def default_test_combinations():
  """Returns the default test combinations for tf.data tests."""
  return combinations.combine(tf_api_version=[1, 2], mode=["eager", "graph"])


def eager_only_combinations():
  """Returns the default test combinations for eager mode only tf.data tests."""
  return combinations.combine(tf_api_version=[1, 2], mode="eager")


def graph_only_combinations():
  """Returns the default test combinations for graph mode only tf.data tests."""
  return combinations.combine(tf_api_version=[1, 2], mode="graph")


def v1_only_combinations():
  """Returns the default test combinations for v1 only tf.data tests."""
  return combinations.combine(tf_api_version=1, mode=["eager", "graph"])


def v2_only_combinations():
  """Returns the default test combinations for v2 only tf.data tests."""
  return combinations.combine(tf_api_version=2, mode=["eager", "graph"])


def v2_eager_only_combinations():
  """Returns the default test combinations for v2 eager only tf.data tests."""
  return combinations.combine(tf_api_version=2, mode="eager")


class DatasetTestBase(test.TestCase):
  """Base class for dataset tests."""

  def assert_op_cancelled(self, op):
    with self.assertRaises(errors.CancelledError):
      self.evaluate(op)

  def assertValuesEqual(self, expected, actual):
    """Asserts that two values are equal."""
    if isinstance(expected, dict):
      self.assertItemsEqual(list(expected.keys()), list(actual.keys()))
      for k in expected.keys():
        self.assertValuesEqual(expected[k], actual[k])
    elif sparse_tensor.is_sparse(expected):
      self.assertAllEqual(expected.indices, actual.indices)
      self.assertAllEqual(expected.values, actual.values)
      self.assertAllEqual(expected.dense_shape, actual.dense_shape)
    else:
      self.assertAllEqual(expected, actual)

  def getNext(self, dataset, requires_initialization=False, shared_name=None):
    """Returns a callable that returns the next element of the dataset.

    Example use:
    ```python
    # In both graph and eager modes
    dataset = ...
    get_next = self.getNext(dataset)
    result = self.evaluate(get_next())
    ```

    Args:
      dataset: A dataset whose elements will be returned.
      requires_initialization: Indicates that when the test is executed in graph
        mode, it should use an initializable iterator to iterate through the
        dataset (e.g. when it contains stateful nodes). Defaults to False.
      shared_name: (Optional.) If non-empty, the returned iterator will be
        shared under the given name across multiple sessions that share the same
        devices (e.g. when using a remote server).
    Returns:
      A callable that returns the next element of `dataset`. Any `TensorArray`
      objects `dataset` outputs are stacked.
    """
    def ta_wrapper(gn):
      def _wrapper():
        r = gn()
        if isinstance(r, tensor_array_ops.TensorArray):
          return r.stack()
        else:
          return r
      return _wrapper

    # Create an anonymous iterator if we are in eager-mode or are graph inside
    # of a tf.function.
    if context.executing_eagerly() or ops.inside_function():
      iterator = iter(dataset)
      return ta_wrapper(iterator._next_internal)  # pylint: disable=protected-access
    else:
      if requires_initialization:
        iterator = dataset_ops.make_initializable_iterator(dataset, shared_name)
        self.evaluate(iterator.initializer)
      else:
        iterator = dataset_ops.make_one_shot_iterator(dataset)
      get_next = iterator.get_next()
      return ta_wrapper(lambda: get_next)

  def _compareOutputToExpected(self, result_values, expected_values,
                               assert_items_equal):
    if assert_items_equal:
      # TODO(shivaniagrawal): add support for nested elements containing sparse
      # tensors when needed.
      self.assertItemsEqual(result_values, expected_values)
      return
    for i in range(len(result_values)):
      nest.assert_same_structure(result_values[i], expected_values[i])
      for result_value, expected_value in zip(
          nest.flatten(result_values[i]), nest.flatten(expected_values[i])):
        self.assertValuesEqual(expected_value, result_value)

  def getDatasetOutput(self, dataset, requires_initialization=False):
    get_next = self.getNext(
        dataset, requires_initialization=requires_initialization)
    return self.getIteratorOutput(get_next)

  def getIteratorOutput(self, get_next):
    """Evaluates `get_next` until end of input, returning the results."""
    results = []
    while True:
      try:
        results.append(self.evaluate(get_next()))
      except errors.OutOfRangeError:
        break
    return results

  def assertDatasetProduces(self,
                            dataset,
                            expected_output=None,
                            expected_shapes=None,
                            expected_error=None,
                            requires_initialization=False,
                            num_test_iterations=1,
                            assert_items_equal=False,
                            expected_error_iter=1):
    """Asserts that a dataset produces the expected output / error.

    Args:
      dataset: A dataset to check for the expected output / error.
      expected_output: A list of elements that the dataset is expected to
        produce.
      expected_shapes: A list of TensorShapes which is expected to match
        output_shapes of dataset.
      expected_error: A tuple `(type, predicate)` identifying the expected error
        `dataset` should raise. The `type` should match the expected exception
        type, while `predicate` should either be 1) a unary function that inputs
        the raised exception and returns a boolean indicator of success or 2) a
        regular expression that is expected to match the error message
        partially.
      requires_initialization: Indicates that when the test is executed in graph
        mode, it should use an initializable iterator to iterate through the
        dataset (e.g. when it contains stateful nodes). Defaults to False.
      num_test_iterations: Number of times `dataset` will be iterated. Defaults
        to 1.
      assert_items_equal: Tests expected_output has (only) the same elements
        regardless of order.
      expected_error_iter: How many times to iterate before expecting an error,
        if an error is expected.
    """
    self.assertTrue(
        expected_error is not None or expected_output is not None,
        "Exactly one of expected_output or expected error should be provided.")
    if expected_error:
      self.assertTrue(
          expected_output is None,
          "Exactly one of expected_output or expected error should be provided."
      )
      with self.assertRaisesWithPredicateMatch(expected_error[0],
                                               expected_error[1]):
        get_next = self.getNext(
            dataset, requires_initialization=requires_initialization)
        for _ in range(expected_error_iter):
          self.evaluate(get_next())
      return
    if expected_shapes:
      self.assertEqual(expected_shapes,
                       dataset_ops.get_legacy_output_shapes(dataset))
    self.assertGreater(num_test_iterations, 0)
    for _ in range(num_test_iterations):
      get_next = self.getNext(
          dataset, requires_initialization=requires_initialization)
      result = []
      for _ in range(len(expected_output)):
        try:
          result.append(self.evaluate(get_next()))
        except errors.OutOfRangeError:
          raise AssertionError(
              "Dataset ended early, producing %d elements out of %d. "
              "Dataset output: %s" %
              (len(result), len(expected_output), str(result)))
      self._compareOutputToExpected(result, expected_output, assert_items_equal)
      with self.assertRaises(errors.OutOfRangeError):
        self.evaluate(get_next())
      with self.assertRaises(errors.OutOfRangeError):
        self.evaluate(get_next())

  def assertDatasetsEqual(self, dataset1, dataset2):
    """Checks that datasets are equal. Supports both graph and eager mode."""
    self.assertTrue(
        structure.are_compatible(
            dataset_ops.get_structure(dataset1),
            dataset_ops.get_structure(dataset2)))

    flattened_types = nest.flatten(
        dataset_ops.get_legacy_output_types(dataset1))

    next1 = self.getNext(dataset1)
    next2 = self.getNext(dataset2)

    while True:
      try:
        op1 = self.evaluate(next1())
      except errors.OutOfRangeError:
        with self.assertRaises(errors.OutOfRangeError):
          self.evaluate(next2())
        break
      op2 = self.evaluate(next2())

      op1 = nest.flatten(op1)
      op2 = nest.flatten(op2)
      assert len(op1) == len(op2)
      for i in range(len(op1)):
        if sparse_tensor.is_sparse(op1[i]) or ragged_tensor.is_ragged(op1[i]):
          self.assertValuesEqual(op1[i], op2[i])
        elif flattened_types[i] == dtypes.string:
          self.assertAllEqual(op1[i], op2[i])
        else:
          self.assertAllClose(op1[i], op2[i])

  def assertDatasetsRaiseSameError(self,
                                   dataset1,
                                   dataset2,
                                   exception_class,
                                   replacements=None):
    """Checks that datasets raise the same error on the first get_next call."""
    if replacements is None:
      replacements = []
    next1 = self.getNext(dataset1)
    next2 = self.getNext(dataset2)
    try:
      self.evaluate(next1())
      raise ValueError(
          "Expected dataset to raise an error of type %s, but it did not." %
          repr(exception_class))
    except exception_class as e:
      expected_message = e.message
      for old, new, count in replacements:
        expected_message = expected_message.replace(old, new, count)
      # Check that the first segment of the error messages are the same.
      with self.assertRaisesRegexp(exception_class,
                                   re.escape(expected_message)):
        self.evaluate(next2())

  def structuredDataset(self, dataset_structure, shape=None,
                        dtype=dtypes.int64):
    """Returns a singleton dataset with the given structure."""
    if shape is None:
      shape = []
    if dataset_structure is None:
      return dataset_ops.Dataset.from_tensors(
          array_ops.zeros(shape, dtype=dtype))
    else:
      return dataset_ops.Dataset.zip(
          tuple([
              self.structuredDataset(substructure, shape, dtype)
              for substructure in dataset_structure
          ]))

  def verifyRandomAccess(self, dataset, expected):
    self.verifyRandomAccessInfiniteCardinality(dataset, expected)
    with self.assertRaises(errors.OutOfRangeError):
      self.evaluate(random_access.at(dataset, index=len(expected)))

  def verifyRandomAccessInfiniteCardinality(self, dataset, expected):
    """Tests randomly accessing elements of a dataset."""
    # Tests accessing the elements in a shuffled order with repeats.
    len_expected = len(expected)
    indices = list(range(len_expected)) * 2
    random.shuffle(indices)
    for i in indices:
      self.assertAllEqual(expected[i],
                          self.evaluate(random_access.at(dataset, i)))

    # Tests accessing the elements in order.
    indices = set(sorted(indices))
    for i in indices:
      self.assertAllEqual(expected[i],
                          self.evaluate(random_access.at(dataset, i)))

  def textFileInitializer(self, vals):
    file = os.path.join(self.get_temp_dir(), "text_file_initializer")
    with open(file, "w") as f:
      f.write("\n".join(str(v) for v in vals) + "\n")
    return lookup_ops.TextFileInitializer(file, dtypes.int64,
                                          lookup_ops.TextFileIndex.LINE_NUMBER,
                                          dtypes.int64,
                                          lookup_ops.TextFileIndex.WHOLE_LINE)

  def keyValueTensorInitializer(self, vals):
    keys_tensor = constant_op.constant(
        list(range(len(vals))), dtype=dtypes.int64)
    vals_tensor = constant_op.constant(vals)
    return lookup_ops.KeyValueTensorInitializer(keys_tensor, vals_tensor)

  def datasetInitializer(self, vals):
    keys = dataset_ops.Dataset.range(len(vals))
    values = dataset_ops.Dataset.from_tensor_slices(vals)
    ds = dataset_ops.Dataset.zip((keys, values))
    return data_lookup_ops.DatasetInitializer(ds)

  def lookupTableInitializer(self, init_source, vals):
    """Returns a lookup table initializer for the given source and values.

    Args:
      init_source: One of ["textfile", "keyvalue", "dataset"], indicating what
        type of initializer to use.
      vals: The initializer values. The keys will be `range(len(vals))`.
    """
    if init_source == "textfile":
      return self.textFileInitializer(vals)
    elif init_source == "keyvaluetensor":
      return self.keyValueTensorInitializer(vals)
    elif init_source == "dataset":
      return self.datasetInitializer(vals)
    else:
      raise ValueError("Unrecognized init_source: " + init_source)

  def graphRoundTrip(self, dataset, allow_stateful=False):
    """Converts a dataset to a graph and back."""
    graph = gen_dataset_ops.dataset_to_graph(
        dataset._variant_tensor, allow_stateful=allow_stateful)  # pylint: disable=protected-access
    return dataset_ops.from_variant(
        gen_experimental_dataset_ops.dataset_from_graph(graph),
        dataset.element_spec)

  def structuredElement(self, element_structure, shape=None,
                        dtype=dtypes.int64):
    """Returns an element with the given structure."""
    if shape is None:
      shape = []
    if element_structure is None:
      return array_ops.zeros(shape, dtype=dtype)
    else:
      return tuple([
          self.structuredElement(substructure, shape, dtype)
          for substructure in element_structure
      ])

  def checkDeterminism(self, dataset_fn, expect_determinism, expected_elements):
    """Tests whether a dataset produces its elements deterministically.

    `dataset_fn` takes a delay_ms argument, which tells it how long to delay
    production of the first dataset element. This gives us a way to trigger
    out-of-order production of dataset elements.

    Args:
      dataset_fn: A function taking a delay_ms argument.
      expect_determinism: Whether to expect deterministic ordering.
      expected_elements: The elements expected to be produced by the dataset,
        assuming the dataset produces elements in deterministic order.
    """
    if expect_determinism:
      dataset = dataset_fn(100)
      actual = self.getDatasetOutput(dataset)
      self.assertAllEqual(expected_elements, actual)
      return

    # We consider the test a success if it succeeds under any delay_ms. The
    # delay_ms needed to observe non-deterministic ordering varies across
    # test machines. Usually 10 or 100 milliseconds is enough, but on slow
    # machines it could take longer.
    for delay_ms in [10, 100, 1000, 20000, 100000]:
      dataset = dataset_fn(delay_ms)
      actual = self.getDatasetOutput(dataset)
      self.assertCountEqual(expected_elements, actual)
      for i in range(len(actual)):
        if actual[i] != expected_elements[i]:
          return
    self.fail("Failed to observe nondeterministic ordering")

  def configureDevicesForMultiDeviceTest(self, num_devices):
    """Configures number of logical devices for multi-device tests.

    It returns a list of device names. If invoked in GPU-enabled runtime, the
    last device name will be for a GPU device. Otherwise, all device names will
    be for a CPU device.

    Args:
      num_devices: The number of devices to configure.

    Returns:
      A list of device names to use for a multi-device test.
    """
    cpus = config.list_physical_devices("CPU")
    gpus = config.list_physical_devices("GPU")
    config.set_logical_device_configuration(cpus[0], [
        context.LogicalDeviceConfiguration() for _ in range(num_devices)
    ])
    devices = ["/device:CPU:" + str(i) for i in range(num_devices - 1)]
    if gpus:
      devices.append("/device:GPU:0")
    else:
      devices.append("/device:CPU:" + str(num_devices - 1))
    return devices