# Copyright 2017 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 for multi-gpu training utilities."""

import tensorflow.compat.v2 as tf

import numpy as np
import keras
from keras import optimizer_v1
from keras.utils import multi_gpu_utils
from keras.utils import np_utils


def check_if_compatible_devices(gpus=2):
  available_devices = [
      keras.utils.multi_gpu_utils._normalize_device_name(name)
      for name in keras.utils.multi_gpu_utils._get_available_devices()
  ]
  if '/gpu:%d' % (gpus - 1) not in available_devices:
    return False
  return True


class TestMultiGPUModel(tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    super(TestMultiGPUModel, self).__init__(methodName)
    gpu_devices = tf.config.list_physical_devices('GPU')
    if len(gpu_devices) == 1:
      # A GPU is available, simulate 2 instead.
      tf.config.set_logical_device_configuration(gpu_devices[0], [
          tf.config.LogicalDeviceConfiguration(500),
          tf.config.LogicalDeviceConfiguration(500)
      ])

  def test_multi_gpu_test_simple_model(self):
    gpus = 2
    num_samples = 1000
    input_dim = 10
    output_dim = 1
    hidden_dim = 10
    epochs = 2
    target_gpu_id = [0, 1]

    if not check_if_compatible_devices(gpus=gpus):
      self.skipTest('multi gpu only')

    with self.cached_session():
      model = keras.models.Sequential()
      model.add(keras.layers.Dense(hidden_dim,
                                   input_shape=(input_dim,)))
      model.add(keras.layers.Dense(output_dim))

      x = np.random.random((num_samples, input_dim))
      y = np.random.random((num_samples, output_dim))

      parallel_model = multi_gpu_utils.multi_gpu_model(model, gpus=gpus)
      parallel_model.compile(loss='mse', optimizer='rmsprop')
      parallel_model.fit(x, y, epochs=epochs)
      parallel_model = multi_gpu_utils.multi_gpu_model(
          model, gpus=target_gpu_id)
      parallel_model.compile(loss='mse', optimizer='rmsprop')
      parallel_model.fit(x, y, epochs=epochs)

  def test_multi_gpu_test_multi_io_model(self):
    gpus = 2
    num_samples = 1000
    input_dim_a = 10
    input_dim_b = 5
    output_dim_a = 1
    output_dim_b = 2
    hidden_dim = 10
    epochs = 2
    target_gpu_id = [0, 1]

    if not check_if_compatible_devices(gpus=gpus):
      self.skipTest('multi gpu only')

    with self.cached_session():
      input_a = keras.Input((input_dim_a,))
      input_b = keras.Input((input_dim_b,))
      a = keras.layers.Dense(hidden_dim)(input_a)
      b = keras.layers.Dense(hidden_dim)(input_b)
      c = keras.layers.concatenate([a, b])
      output_a = keras.layers.Dense(output_dim_a)(c)
      output_b = keras.layers.Dense(output_dim_b)(c)
      model = keras.models.Model([input_a, input_b], [output_a, output_b])

      a_x = np.random.random((num_samples, input_dim_a))
      b_x = np.random.random((num_samples, input_dim_b))
      a_y = np.random.random((num_samples, output_dim_a))
      b_y = np.random.random((num_samples, output_dim_b))

      parallel_model = multi_gpu_utils.multi_gpu_model(model, gpus=gpus)
      parallel_model.compile(loss='mse', optimizer='rmsprop')
      parallel_model.fit([a_x, b_x], [a_y, b_y], epochs=epochs)

      parallel_model = multi_gpu_utils.multi_gpu_model(
          model, gpus=target_gpu_id)
      parallel_model.compile(loss='mse', optimizer='rmsprop')
      parallel_model.fit([a_x, b_x], [a_y, b_y], epochs=epochs)

  def test_multi_gpu_test_invalid_devices(self):
    if not check_if_compatible_devices(gpus=2):
      self.skipTest('multi gpu only')

    with self.cached_session():
      input_shape = (1000, 10)
      model = keras.models.Sequential()
      model.add(keras.layers.Dense(10,
                                   activation='relu',
                                   input_shape=input_shape[1:]))
      model.add(keras.layers.Dense(1, activation='sigmoid'))
      model.compile(loss='mse', optimizer='rmsprop')

      x = np.random.random(input_shape)
      y = np.random.random((input_shape[0], 1))
      with self.assertRaises(ValueError):
        parallel_model = multi_gpu_utils.multi_gpu_model(
            model, gpus=len(keras.backend._get_available_gpus()) + 1)
        parallel_model.fit(x, y, epochs=2)

      with self.assertRaises(ValueError):
        parallel_model = multi_gpu_utils.multi_gpu_model(
            model, gpus=[0, 2, 4, 6, 8])
        parallel_model.fit(x, y, epochs=2)

      with self.assertRaises(ValueError):
        parallel_model = multi_gpu_utils.multi_gpu_model(model, gpus=1)
        parallel_model.fit(x, y, epochs=2)

      with self.assertRaises(ValueError):
        parallel_model = multi_gpu_utils.multi_gpu_model(model, gpus=[0])
        parallel_model.fit(x, y, epochs=2)

  def test_nested_model_with_tensor_input(self):
    gpus = 2
    input_dim = 10
    shape = (input_dim,)
    num_samples = 16
    num_classes = 10

    if not check_if_compatible_devices(gpus=gpus):
      self.skipTest('multi gpu only')

    with tf.Graph().as_default(), self.cached_session():
      input_shape = (num_samples,) + shape
      x_train = np.random.randint(0, 255, input_shape)
      y_train = np.random.randint(0, num_classes, (input_shape[0],))

      y_train = np_utils.to_categorical(y_train, num_classes)

      x_train = x_train.astype('float32')
      y_train = y_train.astype('float32')

      dataset = tf.compat.v1.data.Dataset.from_tensor_slices((x_train, y_train))
      dataset = dataset.repeat()
      dataset = dataset.batch(4)
      iterator = tf.compat.v1.data.make_one_shot_iterator(dataset)

      inputs, targets = iterator.get_next()

      input_tensor = keras.layers.Input(tensor=inputs)

      model = keras.models.Sequential()
      model.add(keras.layers.Dense(3,
                                   input_shape=(input_dim,)))
      model.add(keras.layers.Dense(num_classes))

      output = model(input_tensor)
      outer_model = keras.Model(input_tensor, output)
      parallel_model = multi_gpu_utils.multi_gpu_model(outer_model, gpus=gpus)

      parallel_model.compile(
          loss='categorical_crossentropy',
          optimizer=optimizer_v1.RMSprop(lr=0.0001, decay=1e-6),
          metrics=['accuracy'],
          target_tensors=[targets])
      parallel_model.fit(epochs=1, steps_per_epoch=3)

  def test_multi_gpu_with_multi_input_layers(self):
    gpus = 2

    if not check_if_compatible_devices(gpus=gpus):
      self.skipTest('multi gpu only')

    with self.cached_session():
      inputs = keras.Input((4, 3))
      init_state = keras.Input((3,))
      outputs = keras.layers.SimpleRNN(
          3, return_sequences=True)(inputs, initial_state=init_state)
      x = [np.random.randn(2, 4, 3), np.random.randn(2, 3)]
      y = np.random.randn(2, 4, 3)
      model = keras.Model([inputs, init_state], outputs)
      parallel_model = multi_gpu_utils.multi_gpu_model(model, gpus=gpus)
      parallel_model.compile(loss='mean_squared_error', optimizer='adam')
      parallel_model.train_on_batch(x, y)

  def test_multi_gpu_with_siamese_network(self):
    gpus = 2

    if not check_if_compatible_devices(gpus=gpus):
      self.skipTest('multi gpu only')

    with self.cached_session():
      input_shape = (3,)
      nested_model = keras.models.Sequential([
          keras.layers.Dense(32, input_shape=input_shape),
          keras.layers.Dense(1)
      ], name='nested')

      input1 = keras.Input(input_shape)
      input2 = keras.Input(input_shape)
      score1 = nested_model(input1)
      score2 = nested_model(input2)
      score_sum = keras.layers.Add(name='add')([score1, score2])

      siamese = keras.models.Model(inputs=[input1, input2],
                                   outputs=[score_sum, score1, score2],
                                   name='siamese')
      parallel_siamese = multi_gpu_utils.multi_gpu_model(siamese, gpus)
      self.assertEqual(parallel_siamese.output_names,
                       ['add', 'nested', 'nested_1'])

if __name__ == '__main__':
  tf.test.main()
