# 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.
# ==============================================================================
"""Densenet Training."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from absl import app
import tensorflow as tf # TF2
from tensorflow_examples.models.densenet import densenet
from tensorflow_examples.models.densenet import utils
assert tf.__version__.startswith('2')


class Train(object):
  """Train class.

  Args:
    epochs: Number of epochs
    enable_function: If True, wraps the train_step and test_step in tf.function
    model: Densenet model.
  """

  def __init__(self, epochs, enable_function, model):
    self.epochs = epochs
    self.enable_function = enable_function
    self.autotune = tf.data.experimental.AUTOTUNE
    self.loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
        from_logits=True)
    self.optimizer = tf.keras.optimizers.SGD(learning_rate=0.1,
                                             momentum=0.9, nesterov=True)
    self.train_loss_metric = tf.keras.metrics.Mean(name='train_loss')
    self.train_acc_metric = tf.keras.metrics.SparseCategoricalAccuracy(
        name='train_accuracy')
    self.test_loss_metric = tf.keras.metrics.Mean(name='test_loss')
    self.test_acc_metric = tf.keras.metrics.SparseCategoricalAccuracy(
        name='test_accuracy')
    self.model = model

  def decay(self, epoch):
    if epoch < 150:
      return 0.1
    if epoch >= 150 and epoch < 225:
      return 0.01
    if epoch >= 225:
      return 0.001

  def keras_fit(self, train_dataset, test_dataset):
    self.model.compile(
        optimizer=self.optimizer, loss=self.loss_object, metrics=['accuracy'])
    history = self.model.fit(
        train_dataset, epochs=self.epochs, validation_data=test_dataset,
        verbose=2, callbacks=[tf.keras.callbacks.LearningRateScheduler(
            self.decay)])
    return (history.history['loss'][-1],
            history.history['accuracy'][-1],
            history.history['val_loss'][-1],
            history.history['val_accuracy'][-1])

  def train_step(self, image, label):
    """One train step.

    Args:
      image: Batch of images.
      label: corresponding label for the batch of images.
    """

    with tf.GradientTape() as tape:
      predictions = self.model(image, training=True)
      loss = self.loss_object(label, predictions)
      loss += sum(self.model.losses)
    gradients = tape.gradient(loss, self.model.trainable_variables)
    self.optimizer.apply_gradients(
        zip(gradients, self.model.trainable_variables))

    self.train_loss_metric(loss)
    self.train_acc_metric(label, predictions)

  def test_step(self, image, label):
    """One test step.

    Args:
      image: Batch of images.
      label: corresponding label for the batch of images.
    """

    predictions = self.model(image, training=False)
    loss = self.loss_object(label, predictions)

    self.test_loss_metric(loss)
    self.test_acc_metric(label, predictions)

  def custom_loop(self, train_dataset, test_dataset):
    """Custom training and testing loop.

    Args:
      train_dataset: Training dataset
      test_dataset: Testing dataset

    Returns:
      train_loss, train_accuracy, test_loss, test_accuracy
    """
    if self.enable_function:
      self.train_step = tf.function(self.train_step)
      self.test_step = tf.function(self.test_step)

    for epoch in range(self.epochs):
      self.optimizer.learning_rate = self.decay(epoch)

      for image, label in train_dataset:
        self.train_step(image, label)

      for test_image, test_label in test_dataset:
        self.test_step(test_image, test_label)

      template = ('Epoch: {}, Train Loss: {}, Train Accuracy: {}, '
                  'Test Loss: {}, Test Accuracy: {}')

      print(
          template.format(epoch, self.train_loss_metric.result(),
                          self.train_acc_metric.result(),
                          self.test_loss_metric.result(),
                          self.test_acc_metric.result()))

      if epoch != self.epochs - 1:
        self.train_loss_metric.reset_states()
        self.train_acc_metric.reset_states()
        self.test_loss_metric.reset_states()
        self.test_acc_metric.reset_states()

    return (self.train_loss_metric.result().numpy(),
            self.train_acc_metric.result().numpy(),
            self.test_loss_metric.result().numpy(),
            self.test_acc_metric.result().numpy())


def run_main(argv):
  """Passes the flags to main.

  Args:
    argv: argv
  """
  del argv
  kwargs = utils.flags_dict()
  main(**kwargs)


def main(epochs,
         enable_function,
         buffer_size,
         batch_size,
         mode,
         growth_rate,
         output_classes,
         depth_of_model=None,
         num_of_blocks=None,
         num_layers_in_each_block=None,
         data_format='channels_last',
         bottleneck=True,
         compression=0.5,
         weight_decay=1e-4,
         dropout_rate=0.,
         pool_initial=False,
         include_top=True,
         train_mode='custom_loop',
         data_dir=None):

  model = densenet.DenseNet(mode, growth_rate, output_classes, depth_of_model,
                            num_of_blocks, num_layers_in_each_block,
                            data_format, bottleneck, compression, weight_decay,
                            dropout_rate, pool_initial, include_top)
  train_obj = Train(epochs, enable_function, model)
  train_dataset, test_dataset, _ = utils.create_dataset(
      buffer_size, batch_size, data_format, data_dir)

  print('Training...')
  if train_mode == 'custom_loop':
    return train_obj.custom_loop(train_dataset, test_dataset)
  elif train_mode == 'keras_fit':
    return train_obj.keras_fit(train_dataset, test_dataset)


if __name__ == '__main__':
  utils.define_densenet_flags()
  app.run(run_main)

tensorflow.keras.losses.SparseCategoricalCrossentropy tensorflow.keras.callbacks.LearningRateScheduler tensorflow.GradientTape tensorflow.function tensorflow.keras.metrics.SparseCategoricalAccuracy tensorflow.__version__.startswith tensorflow.keras.metrics.Mean tensorflow.keras.optimizers.SGD
[(25, 'tensorflow.__version__.startswith', 'tf.__version__.startswith', (['"""2"""'], {}), True, 'import tensorflow as tf\n'), (157, 'tensorflow_examples.models.densenet.utils.flags_dict', 'utils.flags_dict', ([], {}), False, 'from tensorflow_examples.models.densenet import utils\n'), (181, 'tensorflow_examples.models.densenet.densenet.DenseNet', 'densenet.DenseNet', (['mode', 'growth_rate', 'output_classes', 'depth_of_model', 'num_of_blocks', 'num_layers_in_each_block', 'data_format', 'bottleneck', 'compression', 'weight_decay', 'dropout_rate', 'pool_initial', 'include_top'], {}), False, 'from tensorflow_examples.models.densenet import densenet\n'), (186, 'tensorflow_examples.models.densenet.utils.create_dataset', 'utils.create_dataset', (['buffer_size', 'batch_size', 'data_format', 'data_dir'], {}), False, 'from tensorflow_examples.models.densenet import utils\n'), (197, 'tensorflow_examples.models.densenet.utils.define_densenet_flags', 'utils.define_densenet_flags', ([], {}), False, 'from tensorflow_examples.models.densenet import utils\n'), (198, 'absl.app.run', 'app.run', (['run_main'], {}), False, 'from absl import app\n'), (41, 'tensorflow.keras.losses.SparseCategoricalCrossentropy', 'tf.keras.losses.SparseCategoricalCrossentropy', ([], {'from_logits': '(True)'}), True, 'import tensorflow as tf\n'), (43, 'tensorflow.keras.optimizers.SGD', 'tf.keras.optimizers.SGD', ([], {'learning_rate': '(0.1)', 'momentum': '(0.9)', 'nesterov': '(True)'}), True, 'import tensorflow as tf\n'), (45, 'tensorflow.keras.metrics.Mean', 'tf.keras.metrics.Mean', ([], {'name': '"""train_loss"""'}), True, 'import tensorflow as tf\n'), (46, 'tensorflow.keras.metrics.SparseCategoricalAccuracy', 'tf.keras.metrics.SparseCategoricalAccuracy', ([], {'name': '"""train_accuracy"""'}), True, 'import tensorflow as tf\n'), (48, 'tensorflow.keras.metrics.Mean', 'tf.keras.metrics.Mean', ([], {'name': '"""test_loss"""'}), True, 'import tensorflow as tf\n'), (49, 'tensorflow.keras.metrics.SparseCategoricalAccuracy', 'tf.keras.metrics.SparseCategoricalAccuracy', ([], {'name': '"""test_accuracy"""'}), True, 'import tensorflow as tf\n'), (81, 'tensorflow.GradientTape', 'tf.GradientTape', ([], {}), True, 'import tensorflow as tf\n'), (117, 'tensorflow.function', 'tf.function', (['self.train_step'], {}), True, 'import tensorflow as tf\n'), (118, 'tensorflow.function', 'tf.function', (['self.test_step'], {}), True, 'import tensorflow as tf\n'), (66, 'tensorflow.keras.callbacks.LearningRateScheduler', 'tf.keras.callbacks.LearningRateScheduler', (['self.decay'], {}), True, 'import tensorflow as tf\n')]