def test_dataset_creator_input_options_with_cluster_coordinator(self):
        dataset_fn = lambda _: dataset_ops.DatasetV2.from_tensor_slices([1, 1])
        input_options = distribute_lib.InputOptions(
            experimental_fetch_to_device=True,
            experimental_per_replica_buffer_size=2)
        x = dataset_creator.DatasetCreator(dataset_fn,
                                           input_options=input_options)
        strategy = self._get_parameter_server_strategy()
        with strategy.scope():
            model = sequential.Sequential([core_layers.Dense(10)])
            model._cluster_coordinator = cluster_coordinator.ClusterCoordinator(
                strategy)
            data_handler = data_adapter.get_data_handler(x,
                                                         steps_per_epoch=2,
                                                         model=model)

        iter_rv = iter(data_handler._dataset)._values[0]
        iter_rv._rebuild_on(model._cluster_coordinator._cluster.workers[0])
        distributed_iterator = iter_rv._get_values()

        # Ensuring the resulting `DistributedIterator` has the right options.
        self.assertTrue(
            distributed_iterator._options.experimental_fetch_to_device)
        self.assertEqual(
            distributed_iterator._options.experimental_per_replica_buffer_size,
            2)
Esempio n. 2
0
  def test_enqueue_incorrect_shape_feature(self):
    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')

    sparse = self._create_high_dimensional_sparse_dataset(strategy)
    sparse_iter = iter(
        strategy.experimental_distribute_dataset(
            sparse,
            options=distribute_lib.InputOptions(
                experimental_fetch_to_device=False)))

    mid_level_api._output_shapes = [TensorShape((1, 1)) for _ in range(3)]
    # The output shape passed to build method is consistent.
    mid_level_api.build([TensorShape([1, 1, 1]) for _ in range(3)])

    @def_function.function
    def test_fn():

      def step():
        return mid_level_api.dequeue()

      mid_level_api.enqueue(next(sparse_iter), training=False)
      return strategy.run(step)

    # Enqueued tensor has shape inconsistent with the output shape setting.
    with self.assertRaisesRegex(ValueError,
                                'Inconsistent shape founded for input feature'):
      test_fn()
Esempio n. 3
0
    def testTypeSpecComponents(self, distribution, enable_get_next_as_optional,
                               experimental_place_dataset_on_device,
                               experimental_fetch_to_device):
        dataset = dataset_ops.DatasetV2.range(10).batch(2)
        distribution.extended.experimental_enable_get_next_as_optional = (
            enable_get_next_as_optional)

        options = distribute_lib.InputOptions(
            experimental_place_dataset_on_device=
            experimental_place_dataset_on_device,
            experimental_fetch_to_device=experimental_fetch_to_device)

        dist_dataset = distribution.experimental_distribute_dataset(
            dataset, options)

        spec = dist_dataset._type_spec
        self.assertEqual(spec._input_workers, dist_dataset._input_workers)
        self.assertEqual(spec._element_spec._value_specs,
                         (tensor_spec.TensorSpec(
                             shape=(None, ), dtype=dtypes.int64, name=None),
                          tensor_spec.TensorSpec(
                              shape=(None, ), dtype=dtypes.int64, name=None)))
        components = spec._to_components(dist_dataset)
        re_dist_dataset = spec._from_components(components)

        self.assertEqual(dist_dataset._input_workers,
                         re_dist_dataset._input_workers)
        self.assertAllEqual(dist_dataset._cloned_datasets,
                            re_dist_dataset._cloned_datasets)
        self.assertEqual(dist_dataset._element_spec,
                         re_dist_dataset._element_spec)
        self.assertEqual(dist_dataset._enable_get_next_as_optional,
                         re_dist_dataset._enable_get_next_as_optional)
        self.assertEqual(dist_dataset._options, re_dist_dataset._options)
Esempio n. 4
0
  def test_enqueue_with_outside_compilation_in_control_flow(self, use_mlir):
    if use_mlir:
      config.enable_mlir_bridge()

    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
    dataset = self._create_sparse_dataset(strategy)
    dataset_iter = iter(strategy.experimental_distribute_dataset(
        dataset,
        options=distribute_lib.InputOptions(
            experimental_prefetch_to_device=False)))

    # This is one way to force the enqueue in some control flow. @tf.functions
    # aren't inlined in the calling tf.function. An alternative would be to
    # place the enqueue in a switch_v2 or something similar.
    @def_function.function
    def enqueue_fn(features):
      mid_level_api.enqueue(features, training=False)

    @def_function.function
    def enqueue_with_outside_compilation():
      def get_activations(features):
        enqueue_fn(features)
        return mid_level_api.dequeue()
      return strategy.run(get_activations, args=(next(dataset_iter),))

    with self.assertRaisesRegex(
        RuntimeError,
        'does not match graph which contains TPUReplicateContext'):
      enqueue_with_outside_compilation()
Esempio n. 5
0
  def test_enqueue_with_outside_compilation_auto_mode(self):
    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
    mid_level_api.build(self.batch_size)
    dataset = self._create_sparse_dataset(strategy)
    dataset_iter = iter(strategy.experimental_distribute_dataset(
        dataset,
        options=distribute_lib.InputOptions(
            experimental_prefetch_to_device=False)))

    @def_function.function
    def enqueue_with_no_gradient_apply(data):
      def get_activations(features):
        # Note the lack of setting training=False, so training defaults to true
        # here even though we don't have apply gradients.
        # We detect the correct mode based on which ops exist that share the
        # same 'name'.
        mid_level_api.enqueue(features, name='call1')
        return mid_level_api.dequeue(name='call1')
      return strategy.run(get_activations, args=(data,))

    @def_function.function
    def enqueue_with_gradient_apply(data):
      def get_activations(features):
        mid_level_api.enqueue(features, name='call2')
        activations = mid_level_api.dequeue(name='call2')
        # Apply an all ones gradient
        gradients = nest.map_structure(array_ops.ones_like, activations)
        mid_level_api.apply_gradients(gradients, name='call2')
        return activations
      return strategy.run(get_activations, args=(data,))

    data = next(dataset_iter)
    before_gradient_apply = enqueue_with_gradient_apply(data)
    after_gradient_apply = enqueue_with_no_gradient_apply(data)
    before_gradient_apply0 = self._get_replica_numpy(before_gradient_apply,
                                                     strategy, 0)
    after_gradient_apply0 = self._get_replica_numpy(after_gradient_apply,
                                                    strategy, 0)

    num_replicas = strategy.num_replicas_in_sync
    # We are passing a gradient of 1 for all lookups, optimizer is SGD with a
    # learning rate of 0.1. Feature 0 and 1 are looked up with a sum combiner
    # with the following ids:
    # Feature 0: [0, 0, 1], [0, 1, 1], ... repeated over num_replicas
    # Feature 1: [0, 1, 1], [0, 0, 1], ... repeated over num_replicas
    # i.e. Row 0 and 1 were looked up 3*num_replicas times over all cores and as
    # the gradient is 1, the accumulated gradient is 3*num_replicas for each
    # position in row 0 and 1 in table.
    #
    # See comments in test_pass_none_to_apply_gradients for the update to
    # Feature 2 and its table.
    # The *2 in the next tests are because those rows have 2 lookups vs
    # the 1 lookup in the other row.
    update = ([[0.3 * num_replicas], [0.3 * num_replicas * 2]],
              [[0.3 * num_replicas * 2], [0.3 * num_replicas]],
              [[0.1 * num_replicas], [0.1 / 3 * num_replicas]])
    golden = tuple([before - np.array(up) for before, up in
                    zip(before_gradient_apply0, update)])

    self.assertAllClose(golden, after_gradient_apply0)
 def test_prefetch_to_device_dataset(self, prefetch_to_device):
     distribution, _, _ = create_test_objects(
         cluster_spec=self._cluster_spec,
         task_type='worker',
         task_id=0,
         num_gpus=2)
     if prefetch_to_device is None:
         input_options = None
     else:
         input_options = distribute_lib.InputOptions(
             experimental_fetch_to_device=prefetch_to_device)
     dataset = dataset_ops.Dataset.range(100)
     dataset = dataset.batch(distribution.num_replicas_in_sync)
     dataset = distribution.experimental_distribute_dataset(  # pylint: disable=assignment-from-no-return
         dataset,
         options=input_options)
     if isinstance(dataset, input_lib.DistributedDatasetV1):
         item = dataset.make_initializable_iterator().get_next()
     else:
         self.skipTest('unsupported test combination')
     device_types = {
         tf_device.DeviceSpec.from_string(tensor.device).device_type
         for tensor in item.values
     }
     self.assertAllEqual(list(device_types), ['GPU'])
  def test_enqueue_with_outside_compilation(self, use_mlir):
    if use_mlir:
      config.enable_mlir_bridge()

    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
    dataset = self._create_sparse_dataset(strategy)
    dataset_iter = iter(strategy.experimental_distribute_dataset(
        dataset,
        options=distribute_lib.InputOptions(
            experimental_prefetch_to_device=False)))

    @def_function.function
    def enqueue_with_outside_compilation(data):
      def get_activations(features):
        mid_level_api.enqueue(features, training=False)
        return mid_level_api.dequeue()
      return strategy.run(get_activations, args=(data,))

    @def_function.function
    def enqueue_without_outside_compilation(data):
      def get_activations():
        return mid_level_api.dequeue()
      mid_level_api.enqueue(data, training=False)
      return strategy.run(get_activations)

    features = next(dataset_iter)

    activations_oc = enqueue_with_outside_compilation(features)
    activations = enqueue_without_outside_compilation(features)

    # Extact per core numpy arrays.
    activations_oc0 = self._get_replica_numpy(activations_oc, strategy, 0)
    activations0 = self._get_replica_numpy(activations, strategy, 0)

    self.assertAllClose(activations_oc0, activations0)
Esempio n. 8
0
    def test_enqueue_with_outside_compilation_non_direct_input(self):
        strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
        mid_level_api.build([
            TensorShape((self.batch_size, 2)),
            TensorShape((self.batch_size, 2)),
            TensorShape((self.batch_size, 3))
        ])
        dataset = self._create_sparse_dataset(strategy)
        dataset_iter = iter(
            strategy.experimental_distribute_dataset(
                dataset,
                options=distribute_lib.InputOptions(
                    experimental_fetch_to_device=False)))

        @def_function.function
        def enqueue_with_outside_compilation():
            def get_activations(features):
                # This inserts a mul operation on the TPU to trigger the direct input
                # error.
                features = (features[0] * 2, features[1] * 2, features[2] * 2)
                mid_level_api.enqueue(features, training=False)
                return mid_level_api.dequeue()

            return strategy.run(get_activations, args=(next(dataset_iter), ))

        with self.assertRaisesRegex(
                ValueError,
                'which does not have the `_tpu_input_identity` attr'):
            enqueue_with_outside_compilation()
Esempio n. 9
0
    def test_sequence_feature(self, is_sparse):
        seq_length = 3
        # Set the max_seq_length in feature config
        for feature in self.feature_config:
            feature.max_sequence_length = seq_length
        strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
        if is_sparse:
            dataset = self._create_sparse_dataset(strategy)
        else:
            dataset = self._create_ragged_dataset(strategy)
        feature_iter = iter(
            strategy.experimental_distribute_dataset(
                dataset,
                options=distribute_lib.InputOptions(
                    experimental_fetch_to_device=False)))

        @def_function.function
        def test_fn():
            def step():
                return mid_level_api.dequeue()

            mid_level_api.enqueue(next(feature_iter), training=False)
            return strategy.run(step)

        output = test_fn()
        self.assertEqual(
            self._get_replica_numpy(output[0], strategy, 0).shape, (2, 3, 4))
        self.assertEqual(
            self._get_replica_numpy(output[1], strategy, 0).shape, (2, 3, 4))
        self.assertEqual(
            self._get_replica_numpy(output[2], strategy, 0).shape, (2, 3, 2))
Esempio n. 10
0
    def test_dense_lookup(self):
        strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')

        input_fn = self._create_dense_input_fn(strategy)
        dist = strategy.distribute_datasets_from_function(
            input_fn,
            options=distribute_lib.InputOptions(
                experimental_prefetch_to_device=False))
        dist_iter = iter(dist)

        @def_function.function
        def test_fn():
            def step():
                return mid_level_api.dequeue()

            mid_level_api.enqueue(next(dist_iter), training=False)
            return strategy.run(step)

        # Run model.
        shard0 = self._get_replica_numpy(test_fn(), strategy, 0)

        # embedding_values is a linear list, so we reshape to match the correct
        # shape of the corresponding table before performing the lookup.
        numpy_videos = np.reshape(self.embedding_values, (8, 4))
        numpy_users = np.reshape(self.embedding_values, (16, 2))
        golden = ((numpy_videos[self.feature_watched_values[-2:]],
                   numpy_videos[self.feature_favorited_values[-2:]],
                   numpy_users[self.feature_friends_values[-2:]]))
        self.assertAllClose(shard0, golden)
    def test_enqueue_wrong_weight_type_for_sparse_and_ragged_tensor(self):
        self.skip_if_oss()
        strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')

        sparse = self._create_sparse_dataset(strategy, include_weights=True)
        ragged = self._create_ragged_dataset(strategy, include_weights=True)
        sparse_iter = iter(
            strategy.experimental_distribute_dataset(
                sparse,
                options=distribute_lib.InputOptions(
                    experimental_fetch_to_device=False)))
        ragged_iter = iter(
            strategy.experimental_distribute_dataset(
                ragged,
                options=distribute_lib.InputOptions(
                    experimental_fetch_to_device=False)))

        @def_function.function
        def test_sparse_fn():
            def step():
                return mid_level_api.dequeue()

            features, _ = next(sparse_iter)
            _, weights = next(ragged_iter)
            mid_level_api.enqueue(features, weights=weights, training=False)
            return strategy.run(step)

        with self.assertRaisesRegex(
                ValueError,
                'which does not match type input which is SparseTensor.'):
            test_sparse_fn()

        @def_function.function
        def test_ragged_fn():
            def step():
                return mid_level_api.dequeue()

            _, weights = next(sparse_iter)
            features, _ = next(ragged_iter)
            mid_level_api.enqueue(features, weights=weights, training=False)
            return strategy.run(step)

        with self.assertRaisesRegex(
                ValueError,
                'which does not match type input which is RaggedTensor.'):
            test_ragged_fn()
Esempio n. 12
0
    def testDoesNotTriggerFunctionTracing(self, distribution,
                                          enable_get_next_as_optional):
        if not tf2.enabled():
            self.skipTest(
                "DistributedIterator CompositeTensor support is only "
                "present in TF 2.0 only.")

        trace_count = [0]

        @def_function.function
        def f(iterator):
            trace_count[0] += 1
            counter = np.int64(0)
            for _ in range(5):
                next(iterator)
                counter += 1
            return counter

        ctx = distribute_lib.InputContext()
        batch_size = ctx.get_per_replica_batch_size(8)
        # Use 20 which isn't divisible by 8 to test partial batch behavior.
        row_lengths = np.mod(np.arange(50), 4).astype(np.int64)
        ragged_tensor = ragged_tensor_lib.RaggedTensor.from_row_lengths(
            np.repeat(np.arange(50, dtype=np.float32), row_lengths),
            row_lengths)
        dataset = dataset_ops.DatasetV2.from_tensor_slices({
            "dense":
            ragged_tensor.to_tensor(),
            "ragged":
            ragged_tensor,
            "sparse":
            ragged_tensor.to_sparse(),
        })
        dataset = dataset.shard(ctx.num_input_pipelines, ctx.input_pipeline_id)
        dataset = dataset.batch(batch_size)

        distribution.extended.experimental_enable_get_next_as_optional = (
            enable_get_next_as_optional)

        if isinstance(distribution,
                      (tpu_strategy.TPUStrategyV2, tpu_strategy.TPUStrategy)):
            # TPUStrategy does not support distributed datasets with device prefetch
            # when using sparse or ragged tensors.
            options = distribute_lib.InputOptions(
                experimental_prefetch_to_device=False)
        else:
            options = None

        dist_dataset = distribution.experimental_distribute_dataset(
            dataset, options)
        with distribution.scope():
            for _ in range(3):
                iterator = iter(dist_dataset)
                _check_type_spec_structure(iterator)
                counter = f(iterator)

                self.assertEqual(trace_count[0], 1)
                self.assertEqual(counter, 5)
Esempio n. 13
0
    def test_pass_none_to_apply_gradients(self):
        strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
        mid_level_api.build([
            TensorShape((self.batch_size, 2)),
            TensorShape((self.batch_size, 2)),
            TensorShape((self.batch_size, 3))
        ])
        dataset = self._create_sparse_dataset(strategy)
        data = next(
            iter(
                strategy.experimental_distribute_dataset(
                    dataset,
                    options=distribute_lib.InputOptions(
                        experimental_fetch_to_device=False))))

        @def_function.function
        def embedding_and_set_gradients(data):
            mid_level_api.enqueue(data)

            def tpu_fn():
                results = mid_level_api.dequeue()
                mid_level_api.apply_gradients(
                    (None, None, array_ops.ones_like(results[2])))
                return results

            return strategy.run(tpu_fn)

        @def_function.function
        def embedding_only(data):
            mid_level_api.enqueue(data, training=False)

            def tpu_fn():
                return mid_level_api.dequeue()

            return strategy.run(tpu_fn)

        first = self._get_replica_numpy(embedding_and_set_gradients(data),
                                        strategy, 0)
        second = self._get_replica_numpy(embedding_only(data), strategy, 0)

        # First two features should be the same as None gradient was applied.
        # Third feature had gradient of 1 passed in from each core.
        # Each core received the same ids per core and returned the following batch:
        # [ row 3, row 0 + row 1 + row 2 ]
        # so gradient update was (learning rate = 0.1):
        #   row 0: -1/3*0.1
        #   row 1: -1/3*0.1
        #   row 2: -1/3*0.1
        #   row 3: -1*0.1
        # There is a factor of num_replicas because each replica gave an update.

        num_replicas = strategy.num_replicas_in_sync
        update = ([[0.0]], [[0.0]], [[0.1 * num_replicas],
                                     [0.1 / 3 * num_replicas]])
        golden = tuple(
            [feature - np.array(up) for feature, up in zip(first, update)])

        self.assertAllClose(golden, second)
    def testFromFunctionInputSignatureForPerReplicaValuesWithOptions(
            self, distribution, enable_get_next_as_optional,
            experimental_place_dataset_on_device,
            experimental_fetch_to_device):

        if experimental_place_dataset_on_device and experimental_fetch_to_device:
            self.skipTest("Setting experimental_place_dataset_on_device and "
                          "experimental_fetch_to_device to `True` is not "
                          "allowed when using "
                          "distribute_lib.InputReplicationMode.PER_REPLICA.")

        fname1 = os.path.join(self.get_temp_dir(), "1.txt")
        _create_text_file(fname1, 5)
        fname2 = os.path.join(self.get_temp_dir(), "2.txt")
        _create_text_file(fname2, 9)

        def dataset_fn(input_context):
            dataset = dataset_ops.DatasetV2.from_tensor_slices(
                [fname1, fname2])
            dataset = dataset.shard(input_context.num_input_pipelines,
                                    input_context.input_pipeline_id)
            return readers.TextLineDatasetV2(dataset).map(
                string_ops.string_to_number).batch(
                    input_context.get_per_replica_batch_size(4))

        options = distribute_lib.InputOptions(
            experimental_place_dataset_on_device=(
                experimental_place_dataset_on_device),
            experimental_fetch_to_device=experimental_fetch_to_device,
            experimental_replication_mode=(
                distribute_lib.InputReplicationMode.PER_REPLICA))

        distribution.extended.experimental_enable_get_next_as_optional = (
            enable_get_next_as_optional)
        ds = distribution.experimental_distribute_datasets_from_function(
            dataset_fn, options)

        iterator = iter(ds)
        _check_type_spec_structure(iterator)
        spec = iterator._type_spec
        tensor_list = spec._to_components(iterator)
        re_iterator = spec._from_components(tensor_list)

        _check_type_spec_structure(iter(ds))
        element_spec = ds.element_spec
        iter_element_spec = iter(ds).element_spec
        nest.assert_same_structure(element_spec, iter_element_spec)
        self.assertAllEqual(nest.flatten(element_spec),
                            nest.flatten(iter_element_spec))
        self.assertEqual(iterator._input_workers, re_iterator._input_workers)
        self.assertAllEqual(iterator._iterators, re_iterator._iterators)

        @def_function.function(input_signature=[element_spec])
        def process_inputs(inputs):
            distribution.run(lambda inputs: inputs, args=(inputs, ))

        for x in ds:
            process_inputs(x)
Esempio n. 15
0
  def test_enqueue_dense_sparse_ragged(self):
    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')

    dataset = self._create_high_dimensional_dense_dataset(strategy)
    dense_iter = iter(
        strategy.experimental_distribute_dataset(
            dataset,
            options=distribute_lib.InputOptions(
                experimental_fetch_to_device=False)))

    sparse = self._create_high_dimensional_sparse_dataset(strategy)
    sparse_iter = iter(
        strategy.experimental_distribute_dataset(
            sparse,
            options=distribute_lib.InputOptions(
                experimental_fetch_to_device=False)))

    ragged = self._create_high_dimensional_ragged_dataset(strategy)
    ragged_iter = iter(
        strategy.experimental_distribute_dataset(
            ragged,
            options=distribute_lib.InputOptions(
                experimental_fetch_to_device=False)))

    mid_level_api.build([
        TensorShape([self.batch_size, self.data_batch_size, 1]),
        TensorShape([self.batch_size, self.data_batch_size, 2]),
        TensorShape([self.batch_size, self.data_batch_size, 3])
    ])

    @def_function.function
    def test_fn():

      def step():
        return mid_level_api.dequeue()

      features = (next(dense_iter)[0], next(sparse_iter)[1],
                  next(ragged_iter)[2])
      mid_level_api.enqueue(features, training=False)
      return strategy.run(step)

    test_fn()
Esempio n. 16
0
    def test_enqueue_with_weights(self, ragged):
        strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
        weight = 0.5
        if ragged:
            dataset = self._create_ragged_dataset(strategy,
                                                  include_weights=True,
                                                  weight=weight)
        else:
            dataset = self._create_sparse_dataset(strategy,
                                                  include_weights=True,
                                                  weight=weight)
            mid_level_api.build([
                TensorShape((self.batch_size, 2)),
                TensorShape((self.batch_size, 2)),
                TensorShape((self.batch_size, 3))
            ])

        dataset_iter = iter(
            strategy.experimental_distribute_dataset(
                dataset,
                options=distribute_lib.InputOptions(
                    experimental_fetch_to_device=False)))

        @def_function.function
        def enqueue_and_get(features, weights):
            def get_activations():
                return mid_level_api.dequeue()

            mid_level_api.enqueue(features, weights=weights, training=False)
            return strategy.run(get_activations)

        features, weights = next(dataset_iter)
        # Replace the weight for the second feature by None to test.
        weights = (weights[0], None, weights[2])

        no_weights_activations = enqueue_and_get(features, weights=None)
        weights_activations = enqueue_and_get(features, weights=weights)

        # Extact per core numpy arrays.
        no_weights0 = self._get_replica_numpy(no_weights_activations, strategy,
                                              0)
        weights0 = self._get_replica_numpy(weights_activations, strategy, 0)
        # videos table has sum combiner and users table has mean combiner.
        # i.e. users table lookups isn't affected by the weights as all the weights
        # are the same.
        # Tuple entry 0 and 1 are the watched and favorited features from the videos
        # table and entry 2 is the friends feature from the users table.
        # Note that None was passed as a weight for entry 1 so weight should have no
        # effect.
        weight = (0.5, 1.0, 1.0)
        golden = tuple(
            [no_weight * w for no_weight, w in zip(no_weights0, weight)])

        self.assertAllClose(golden, weights0)
Esempio n. 17
0
  def test_dataset_creator_model_fit_without_strategy(self, use_input_options):
    model = sequential.Sequential([core_layers.Dense(10)])
    model.compile(gradient_descent.SGD(), loss="mse")

    input_options = distribute_lib.InputOptions() if use_input_options else None
    history = model.fit(
        dataset_creator.DatasetCreator(self._get_dataset_fn(), input_options),
        epochs=10,
        steps_per_epoch=10,
        verbose=0)
    self.assertLen(history.history["loss"], 10)
Esempio n. 18
0
  def test_composite_input_non_flat_output(self, enable_packed_var):
    strategy = get_tpu_strategy(enable_packed_var)
    if strategy.num_replicas_in_sync != 2:
      self.skipTest("Test assumes two replicas.")

    with strategy.scope():
      table = variables.Variable(
          initial_value=[[0.0, 1.0], [3.0, 7.0]], dtype=dtypes.float32)

    @def_function.function
    def sparse_lookup(iterator):

      def tpu_function(sparse):
        # Assumes dense_shape is (2, *)
        looked_up = array_ops.gather(table, sparse.values)
        segment_sum = math_ops.unsorted_segment_sum(
            looked_up, sparse.indices[:, 0], 2)
        return {"sparse": sparse, "segment_sum": segment_sum}

      return nest.map_structure(
          strategy.experimental_local_results,
          strategy.run(tpu_function, args=(next(iterator),)))

    def dataset_fn(_):
      dataset = dataset_ops.Dataset.range(2)

      def make_sparse(_):
        return sparse_tensor.SparseTensor(
            indices=array_ops.constant([[0, 0], [1, 0], [1, 1]],
                                       dtype=dtypes.int64),
            values=array_ops.constant([0, 0, 1], dtype=dtypes.int32),
            dense_shape=array_ops.constant([2, 2], dtype=dtypes.int64))

      return dataset.map(make_sparse)

    dataset = iter(
        strategy.experimental_distribute_datasets_from_function(
            dataset_fn,
            distribute_lib.InputOptions(
                experimental_prefetch_to_device=False)))

    output = sparse_lookup(dataset)

    # All replicas return identical reults.
    for replica in range(strategy.num_replicas_in_sync):
      self.assertIsInstance(output["sparse"][replica],
                            sparse_tensor.SparseTensor)
      self.assertAllEqual(output["sparse"][replica].indices,
                          [[0, 0], [1, 0], [1, 1]])
      self.assertAllEqual(output["sparse"][replica].values, [0, 0, 1])
      self.assertAllEqual(output["sparse"][replica].dense_shape, [2, 2])
      self.assertAllEqual(output["segment_sum"][replica],
                          [[0.0, 1.0], [3.0, 8.0]])
  def test_embedding(self, optimizer_name, training, sparse):
    strategy, mid_level_api, optimizer = (
        self._create_strategy_and_mid_level(optimizer_name))

    if sparse:
      dataset = self._create_sparse_dataset(strategy)
    else:
      dataset = self._create_ragged_dataset(strategy)

    dist = strategy.experimental_distribute_dataset(
        dataset,
        options=distribute_lib.InputOptions(experimental_fetch_to_device=False))
    dist_iter = iter(dist)

    @def_function.function
    def test_fn():

      def step():
        """Create and run computation that returns the embedding activations."""
        if not training:
          activations = mid_level_api.dequeue()
          total_loss = _get_total_loss_tensor(activations)
          ret_val = [total_loss] + list(activations)
          return ret_val
        else:
          with backprop.GradientTape() as tape:
            activations = mid_level_api.dequeue()
            tape.watch(activations)
            total_loss = _get_total_loss_tensor(activations)
            loss_per_replica = total_loss / strategy.num_replicas_in_sync
          gradients = tape.gradient(loss_per_replica, activations)
          mid_level_api.apply_gradients(gradients)
        ret_val = [total_loss] + list(activations)
        return ret_val

      mid_level_api.enqueue(next(dist_iter), training=training)
      result = strategy.run(step)
      return result

    # Run model.
    shard_out_val = test_fn()

    # Retrieve TPU weights to CPU.
    mid_level_api._retrieve_variables()

    # Compute sparse tensors for global batch.
    input_data = next(iter(self._create_sparse_dataset(strategy)))

    # Check results.
    self._check_results(strategy, shard_out_val, training, input_data,
                        mid_level_api._variables,
                        optimizer)
Esempio n. 20
0
  def test_prefetch_to_device_tpu(self):
    strategy = get_tpu_strategy()
    dataset = dataset_ops.Dataset.range(
        strategy.num_replicas_in_sync * 2,
        output_type=dtypes.float32).batch(strategy.num_replicas_in_sync)

    input_options = distribute_lib.InputOptions(
        experimental_fetch_to_device=True)
    dataset_item = next(iter(strategy.experimental_distribute_dataset(
        dataset, options=input_options)))
    dataset_location = tf_device.DeviceSpec.from_string(
        dataset_item.values[0].device)
    self.assertEqual(dataset_location.device_type, "TPU")
Esempio n. 21
0
  def testDistributeDatasetFunctionHostPrefetch(self, distribution):
    data = [5., 6., 7., 8.]
    input_iterator = iter(
        distribution.distribute_datasets_from_function(
            lambda _: get_dataset_from_tensor_slices(data),
            distribute_lib.InputOptions(experimental_fetch_to_device=False)))

    local_results = distribution.experimental_local_results(
        input_iterator.get_next())

    for result in local_results:
      self.assertEqual(result.backing_device,
                       device_util.resolve("/device:CPU:0"))
Esempio n. 22
0
  def test_different_input_shapes(self):
    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')

    sparse = self._create_high_dimensional_sparse_dataset(strategy)
    sparse_iter = iter(
        strategy.experimental_distribute_dataset(
            sparse,
            options=distribute_lib.InputOptions(
                experimental_fetch_to_device=False)))
    # Create a feature with shape (1, 3, 1)
    dense_feature = constant_op.constant(
        np.zeros(3), shape=(1, 3, 1), dtype=dtypes.int32)
    dense_dataset = dataset_ops.DatasetV2.from_tensors(
        dense_feature).unbatch().repeat().batch(
            1 * strategy.num_replicas_in_sync, drop_remainder=True)
    dense_iter = iter(
        strategy.experimental_distribute_dataset(
            dense_dataset,
            options=distribute_lib.InputOptions(
                experimental_fetch_to_device=False)))

    @def_function.function
    def test_fn():

      def step():
        return mid_level_api.dequeue()

      features = (next(dense_iter), next(sparse_iter)[1], next(sparse_iter)[2])
      mid_level_api.enqueue(features, training=False)
      return strategy.run(step)

    test_fn()

    self.assertEqual(mid_level_api._output_shapes, [
        TensorShape((1, 3)),
        TensorShape((self.batch_size, self.data_batch_size)),
        TensorShape((self.batch_size, self.data_batch_size))
    ])
    def test_enqueue_with_weights(self, ragged):
        strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
        weight = 0.5
        if ragged:
            dataset = self._create_ragged_dataset(strategy,
                                                  include_weights=True,
                                                  weight=weight)
        else:
            dataset = self._create_sparse_dataset(strategy,
                                                  include_weights=True,
                                                  weight=weight)

        dataset_iter = iter(
            strategy.experimental_distribute_dataset(
                dataset,
                options=distribute_lib.InputOptions(
                    experimental_fetch_to_device=False)))

        @def_function.function
        def embedding_lookup(features, weights):
            def step(features, weights):
                return mid_level_api(features, weights)

            return strategy.run(step, args=(features, weights))

        features, weights = next(dataset_iter)
        # Replace the weight for the second feature by None to test.
        weights = (weights[0], None, weights[2])

        no_weights_activations = embedding_lookup(features, weights=None)
        weights_activations = embedding_lookup(features, weights=weights)

        no_weights0 = (self._unpack(strategy, no_weights_activations[0]),
                       self._unpack(strategy, no_weights_activations[1]),
                       self._unpack(strategy, no_weights_activations[2]))
        weights0 = (self._unpack(strategy, weights_activations[0]),
                    self._unpack(strategy, weights_activations[1]),
                    self._unpack(strategy, weights_activations[2]))
        # videos table has sum combiner and users table has mean combiner.
        # i.e. users table lookups isn't affected by the weights as all the weights
        # are the same.
        # Tuple entry 0 and 1 are the watched and favorited features from the videos
        # table and entry 2 is the friends feature from the users table.
        # Note that None was passed as a weight for entry 1 so weight should have no
        # effect.
        weight = (0.5, 1.0, 1.0)
        golden = tuple(
            [no_weight * w for no_weight, w in zip(no_weights0, weight)])

        self.assertAllClose(golden, weights0)
Esempio n. 24
0
    def testTypeSpecRoundTrip(self, distribution, enable_get_next_as_optional):
        if not tf2.enabled():
            self.skipTest(
                "DistributedIterator CompositeTensor support is only "
                "present in TF 2.0 only.")

        ctx = distribute_lib.InputContext()
        batch_size = ctx.get_per_replica_batch_size(8)
        # Use 20 which isn't divisible by 8 to test partial batch behavior.
        row_lengths = np.mod(np.arange(20), 4).astype(np.int64)
        ragged_tensor = ragged_tensor_lib.RaggedTensor.from_row_lengths(
            np.repeat(np.arange(20, dtype=np.float32), row_lengths),
            row_lengths)
        dataset = dataset_ops.DatasetV2.from_tensor_slices({
            "dense":
            ragged_tensor.to_tensor(),
            "ragged":
            ragged_tensor,
            "sparse":
            ragged_tensor.to_sparse(),
        })
        dataset = dataset.shard(ctx.num_input_pipelines, ctx.input_pipeline_id)
        dataset = dataset.batch(batch_size)

        distribution.extended.experimental_enable_get_next_as_optional = (
            enable_get_next_as_optional)

        if isinstance(distribution,
                      (tpu_strategy.TPUStrategyV2, tpu_strategy.TPUStrategy)):
            # TPUStrategy does not support distributed datasets with device prefetch
            # when using sparse or ragged tensors.
            options = distribute_lib.InputOptions(
                experimental_prefetch_to_device=False)
        else:
            options = None

        dist_dataset = distribution.experimental_distribute_dataset(
            dataset, options)
        with distribution.scope():
            iterator = iter(dist_dataset)
            _check_type_spec_structure(iterator)

        spec = iterator._type_spec

        tensor_list = spec._to_components(iterator)
        re_iterator = spec._from_components(tensor_list)

        self.assertEqual(iterator._input_workers, re_iterator._input_workers)
        self.assertAllEqual(iterator._iterators, re_iterator._iterators)
Esempio n. 25
0
 def test_prefetch_to_host_dataset(self, distribution):
     input_options = distribute_lib.InputOptions(
         experimental_fetch_to_device=False)
     dataset = dataset_ops.Dataset.range(100)
     dataset = dataset.batch(distribution.num_replicas_in_sync)
     dataset = distribution.experimental_distribute_dataset(
         dataset, options=input_options)
     if context.executing_eagerly():
         item = next(iter(dataset))
     else:
         if isinstance(dataset, input_lib_v1.DistributedDatasetV1):
             item = dataset.make_initializable_iterator().get_next()
         else:
             self.skipTest("unsupported test combination")
     self.assertAllEqual(
         tf_device.DeviceSpec.from_string(item.device).device_type, "CPU")
Esempio n. 26
0
    def test_enqueue_per_device(self):
        self.skip_if_oss()
        strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')

        sparse = self._create_sparse_dataset(strategy)
        sparse_iter = iter(
            strategy.experimental_distribute_dataset(
                sparse,
                options=distribute_lib.InputOptions(
                    experimental_fetch_to_device=False)))

        @def_function.function
        def test_fn():
            def get_activations(dense_value):
                return mid_level_api.dequeue(), dense_value

            sparse_features = next(sparse_iter)
            mid_level_api.enqueue(sparse_features, training=False)
            activations, dense_value1 = strategy.run(get_activations,
                                                     args=(0.0, ))

            def enqueue_fn(ctx):
                core_id = ctx.replica_id_in_sync_group
                device = strategy.extended.worker_devices[core_id]
                sparse_features_local = nest.map_structure(
                    lambda x: strategy.experimental_local_results(x)[core_id],
                    sparse_features)
                mid_level_api.enqueue(sparse_features_local,
                                      training=False,
                                      device=device)
                return 0.0

            data = strategy.experimental_distribute_values_from_function(
                enqueue_fn)
            per_device_activations, dense_value2 = strategy.run(
                get_activations, args=(data, ))
            return activations, per_device_activations, dense_value1, dense_value2

        activations, per_device_activations, _, _ = test_fn()

        # Extact per core numpy arrays and check that both sparse and ragged have
        # the same results.
        activations0 = self._get_replica_numpy(activations, strategy, 0)
        per_device_activations0 = self._get_replica_numpy(
            per_device_activations, strategy, 0)
        self.assertAllClose(activations0, per_device_activations0)
        test_fn()
Esempio n. 27
0
  def test_dataset_creator_input_options(self):
    dataset_fn = lambda _: dataset_ops.DatasetV2.from_tensor_slices([1, 1])
    input_options = distribute_lib.InputOptions(
        experimental_fetch_to_device=True,
        experimental_per_replica_buffer_size=2)
    x = dataset_creator.DatasetCreator(dataset_fn, input_options=input_options)
    with collective_all_reduce_strategy.CollectiveAllReduceStrategy().scope():
      data_handler = data_adapter.get_data_handler(
          x,
          steps_per_epoch=2,
          model=sequential.Sequential([core_layers.Dense(10)]))

    # Ensuring the resulting `DistributedDatasetsFromFunction` has the right
    # options.
    self.assertTrue(data_handler._dataset._options.experimental_fetch_to_device)
    self.assertEqual(
        data_handler._dataset._options.experimental_per_replica_buffer_size, 2)
Esempio n. 28
0
    def test_composite_input_dynamic_shapes_outside_compilation(
            self, enable_packed_var):
        strategy = get_tpu_strategy(enable_packed_var)
        if strategy.num_replicas_in_sync != 2:
            self.skipTest("Test assumes two replicas.")

        table = variables.Variable(initial_value=[[0.0, 1.0], [3.0, 7.0]],
                                   dtype=dtypes.float32)

        @def_function.function
        def sparse_lookup(iterator):
            def tpu_function(sparse):
                lookup = tpu.outside_compilation(
                    embedding_ops.safe_embedding_lookup_sparse, table, sparse)
                return math_ops.reduce_sum(lookup, axis=0)

            return strategy.experimental_local_results(
                strategy.run(tpu_function, args=(next(iterator), )))

        def dataset_fn(_):
            dataset = dataset_ops.Dataset.range(2)

            def make_sparse(i):
                indices = array_ops.constant([[0, 0], [1, 0], [1, 1]],
                                             dtype=dtypes.int64)[0:2 + i]
                values = array_ops.constant([0, 0, 1],
                                            dtype=dtypes.int32)[0:2 + i]
                shape = [
                    array_ops.constant([2], dtype=dtypes.int64),
                    array_ops.expand_dims(1 + i, axis=0)
                ]
                dense_shape = array_ops.concat(shape, axis=0)
                return sparse_tensor.SparseTensor(indices=indices,
                                                  values=values,
                                                  dense_shape=dense_shape)

            return dataset.map(make_sparse)

        dataset = iter(
            strategy.experimental_distribute_datasets_from_function(
                dataset_fn,
                options=distribute_lib.InputOptions(
                    experimental_prefetch_to_device=False)))

        result = sparse_lookup(dataset)
        self.assertAllEqual(result, [[0.0, 2.0], [1.5, 5.0]])
Esempio n. 29
0
 def test_prefetch_to_host_dataset(self):
     distribution, _, _ = self._get_test_object(task_type='worker',
                                                task_id=0,
                                                num_gpus=2)
     input_options = distribute_lib.InputOptions(
         experimental_fetch_to_device=False)
     dataset = dataset_ops.Dataset.range(100)
     dataset = dataset.batch(distribution.num_replicas_in_sync)
     dataset = distribution.experimental_distribute_dataset(
         dataset, options=input_options)
     if isinstance(dataset, input_lib.DistributedDatasetV1):
         item = dataset.make_initializable_iterator().get_next()
     else:
         self.skipTest('unsupported test combination')
     device_types = {
         tf_device.DeviceSpec.from_string(tensor.device).device_type
         for tensor in item.values
     }
     self.assertAllEqual(list(device_types), ['CPU'])
Esempio n. 30
0
  def test_enqueue_weight_for_dense_tensor(self):
    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')

    input_fn = self._create_dense_input_fn(strategy, include_weights=True)
    dist = strategy.distribute_datasets_from_function(
        input_fn,
        options=distribute_lib.InputOptions(
            experimental_prefetch_to_device=False))
    dist_iter = iter(dist)

    @def_function.function
    def test_fn():
      def step():
        return mid_level_api.dequeue()

      features, weights = next(dist_iter)
      mid_level_api.enqueue(features, weights=weights, training=False)
      return strategy.run(step)

    with self.assertRaisesRegex(ValueError, 'Weight specified for dense input'):
      test_fn()