text: # 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.
  # ==============================================================================
  """Tests that show that DistributionStrategy works with canned Estimator."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import shutil
  import tempfile
  from absl.testing import parameterized
  import numpy as np
  import six
  
  from tensorflow.contrib.distribute.python import combinations
  from tensorflow.contrib.distribute.python import mirrored_strategy
  from tensorflow.core.protobuf import config_pb2
  from tensorflow.python import keras
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.eager import context
  from tensorflow.python.estimator import run_config
  from tensorflow.python.estimator import training
  from tensorflow.python.estimator.canned import dnn_linear_combined
  from tensorflow.python.estimator.canned import prediction_keys
  from tensorflow.python.estimator.export import export
  from tensorflow.python.estimator.inputs import numpy_io
  from tensorflow.python.feature_column import feature_column
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.keras.optimizer_v2 import adam
  from tensorflow.python.keras.optimizer_v2 import gradient_descent
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import variable_scope
  from tensorflow.python.ops import variables
  from tensorflow.python.platform import gfile
  from tensorflow.python.platform import test
  from tensorflow.python.summary.writer import writer_cache
  from tensorflow.python.training import distribution_strategy_context as ds_context
  
  
  class KerasOptimizerV2IntegrationTest(test.TestCase, parameterized.TestCase):
  
    def setUp(self):
      self._model_dir = tempfile.mkdtemp()
  
    def dataset_input_fn(self, x, y, batch_size):
  
      def input_fn():
        dataset = dataset_ops.Dataset.from_tensor_slices((x, y))
        dataset = dataset.repeat(1).batch(batch_size)
        return dataset
  
      return input_fn
  
    @combinations.generate(
        combinations.combine(
            mode=['graph'],
            distribution=[
                combinations.one_device_strategy,
                combinations.mirrored_strategy_with_gpu_and_cpu,
                combinations.mirrored_strategy_with_two_gpus
            ],
            use_train_and_evaluate=[True, False]))
    def test_complete_flow_with_mode(self, distribution, use_train_and_evaluate):
      label_dimension = 2
      input_dimension = label_dimension
      batch_size = 10
      data = np.linspace(0., 2., batch_size * label_dimension, dtype=np.float32)
      data = data.reshape(batch_size, label_dimension)
      train_input_fn = self.dataset_input_fn(
          x={'x': data},
          y=data,
          batch_size=batch_size // len(distribution.worker_devices))
      eval_input_fn = self.dataset_input_fn(
          x={'x': data},
          y=data,
          batch_size=batch_size // len(distribution.worker_devices))
      predict_input_fn = numpy_io.numpy_input_fn(
          x={'x': data}, batch_size=batch_size, shuffle=False)
  
      linear_feature_columns = [
          feature_column.numeric_column('x', shape=(input_dimension,))
      ]
      dnn_feature_columns = [
          feature_column.numeric_column('x', shape=(input_dimension,))
      ]
      feature_columns = linear_feature_columns + dnn_feature_columns
      session_config = config_pb2.ConfigProto(
          log_device_placement=True, allow_soft_placement=True)
      estimator = dnn_linear_combined.DNNLinearCombinedRegressor(
          linear_feature_columns=linear_feature_columns,
          dnn_hidden_units=(2, 2),
          dnn_feature_columns=dnn_feature_columns,
          label_dimension=label_dimension,
          model_dir=self._model_dir,
          dnn_optimizer=adam.Adam(0.001),
          linear_optimizer=adam.Adam(0.001),
          config=run_config.RunConfig(
              train_distribute=distribution,
              eval_distribute=distribution,
              session_config=session_config))
  
      num_steps = 2
      if use_train_and_evaluate:
        scores, _ = training.train_and_evaluate(
            estimator, training.TrainSpec(train_input_fn, max_steps=num_steps),
            training.EvalSpec(eval_input_fn))
      else:
        estimator.train(train_input_fn, steps=num_steps)
        scores = estimator.evaluate(eval_input_fn)
  
      self.assertIn('loss', six.iterkeys(scores))
  
      predictions = np.array([
          x[prediction_keys.PredictionKeys.PREDICTIONS]
          for x in estimator.predict(predict_input_fn)
      ])
      self.assertAllEqual((batch_size, label_dimension), predictions.shape)
  
      feature_spec = feature_column.make_parse_example_spec(feature_columns)
      serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn(
          feature_spec)
      export_dir = estimator.export_savedmodel(tempfile.mkdtemp(),
                                               serving_input_receiver_fn)
      self.assertTrue(gfile.Exists(export_dir))
  
    def tearDown(self):
      if self._model_dir:
        writer_cache.FileWriterCache.clear()
        shutil.rmtree(self._model_dir)
  
  
  def get_model():
    x = keras.layers.Input(shape=(3,), name='input')
    y = keras.layers.Dense(4, name='dense')(x)
    model = keras.Model(x, y)
    return model
  
  
  class MirroredStrategyOptimizerV2Test(test.TestCase):
  
    def testKerasOptimizerWithUnequalInput(self):
      if context.num_gpus() < 1:
        self.skipTest('Not enough GPUs.')
  
      def create_fn():
        var = variables.Variable(
            2.0, name='var', aggregation=variable_scope.VariableAggregation.SUM)
        # grad for cpu is 1, grad for gpu is 2, avg grad is 1.5.
        loss = math_ops.cast(_replica_id() + 1, dtype=dtypes.float32) * var
        optimizer = adam.Adam(learning_rate=0.01, beta_1=0.2, beta_2=0.2)
        train_op = optimizer.minimize(loss, var_list=[var])
        m = optimizer.get_slot(var, 'm')
        v = optimizer.get_slot(var, 'v')
        return (var, m, v, train_op, optimizer.iterations)
  
      devices = ['/device:GPU:0', '/device:CPU:0']
      dist = mirrored_strategy.MirroredStrategy(devices)
      with dist.scope():
        (var, m, v, op, counter) = dist.call_for_each_replica(create_fn)
        self.evaluate(variables.global_variables_initializer())
        var_val = [2.0, 2.0, 2.0]
        self.assertAllClose(
            var_val,
            self.evaluate(
                [dist.read_var(var),
                 var.get(devices[0]),
                 var.get(devices[1])]))
        self.assertAllClose([0, 0, 0],
                            self.evaluate([
                                dist.read_var(counter),
                                counter.get(devices[0]),
                                counter.get(devices[1])
                            ]))
  
        train_op = dist.unwrap(op)
        self.evaluate(train_op)
        # m(1) = beta1 * m(0) + (1-beta1) * grad = 0.2 * 0 + 0.8 * (1 + 2) / 2
        m_val = [1.2, 1.2, 1.2]
        # assert slot variables in both replicas are the same.
        self.assertAllClose(
            m_val,
            self.evaluate(
                [dist.read_var(m),
                 m.get(devices[0]),
                 m.get(devices[1])]))
        # v(1) = beta2 * v(0) + (1-beta2) * grad^2 = 0.2 * 0 + 0.8 * 2.25
        v_val = [1.8, 1.8, 1.8]
        self.assertAllClose(
            v_val,
            self.evaluate(
                [dist.read_var(v),
                 v.get(devices[0]),
                 v.get(devices[1])]))
        # var(1) = var(0) - lr * m(1) * sqrt(1 - beta2) / sqrt(v(1)) / (1 - beta1)
        #        = 2.0 - 0.01 * 1.2 * sqrt(0.8) / sqrt(1.8) / 0.8
        var_val = [1.99, 1.99, 1.99]
        self.assertAllClose(
            var_val,
            self.evaluate(
                [dist.read_var(var),
                 var.get(devices[0]),
                 var.get(devices[1])]))
        self.assertAllClose([1, 1, 1],
                            self.evaluate([
                                dist.read_var(counter),
                                counter.get(devices[0]),
                                counter.get(devices[1])
                            ]))
  
        self.evaluate(train_op)
        # m(2) = beta1 * m(1) + (1-beta1) * grad = 0.2 * 1.2 + 0.8 * 1.5
        m_val = [1.44, 1.44, 1.44]
        self.assertAllClose(
            m_val,
            self.evaluate(
                [dist.read_var(m),
                 m.get(devices[0]),
                 m.get(devices[1])]))
        # v(2) = beta2 * v(1) + (1-beta2) * grad^2 = 0.2 * 1.8 + 0.8 * 2.25
        v_val = [2.16, 2.16, 2.16]
        self.assertAllClose(
            v_val,
            self.evaluate(
                [dist.read_var(v),
                 v.get(devices[0]),
                 v.get(devices[1])]))
        self.assertAllClose([2, 2, 2],
                            self.evaluate([
                                dist.read_var(counter),
                                counter.get(devices[0]),
                                counter.get(devices[1])
                            ]))
  
    def testOptimizerWithKerasModelAndNumpyArrays(self):
      if context.num_gpus() < 1:
        self.skipTest('Not enough GPUs.')
  
      with self.cached_session():
        model = get_model()
        optimizer = gradient_descent.SGD(0.001)
        loss = 'mse'
        metrics = ['mae']
        devices = ['/device:GPU:0', '/device:CPU:0']
        dist = mirrored_strategy.MirroredStrategy(devices)
        model.compile(optimizer, loss, metrics=metrics, distribute=dist)
  
        inputs = np.zeros((64, 3), dtype=np.float32)
        targets = np.zeros((64, 4), dtype=np.float32)
  
        model.fit(
            inputs,
            targets,
            epochs=1,
            batch_size=2,
            verbose=0,
            validation_data=(inputs, targets))
        model.evaluate(inputs, targets)
        model.predict(inputs)
  
  
  def _replica_id():
    # TODO(cjfj): Return `replica_id_...` directly, once it is a `Tensor`.
    return constant_op.constant(
        ds_context.get_replica_context().replica_id_in_sync_group)
  
  
  if __name__ == '__main__':
    test.main()
  