# coding=utf-8
# Copyright 2021 The Google Research Authors.
#
# 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.

"""ColTran: Training and Continuous Evaluation."""

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

import collections
import functools
import os
import time

from absl import app
from absl import flags
from absl import logging

from ml_collections import config_flags
import tensorflow as tf
import tensorflow_datasets as tfds

from coltran import datasets
from coltran.models import colorizer
from coltran.models import upsampler
from coltran.utils import train_utils


# pylint: disable=g-direct-tensorflow-import

# pylint: disable=missing-docstring
# pylint: disable=not-callable
# pylint: disable=g-long-lambda

flags.DEFINE_enum('mode', 'train', [
    'train', 'eval_train', 'eval_valid', 'eval_test'], 'Operation mode.')

flags.DEFINE_string('logdir', '/tmp/svt', 'Main directory for logs.')
flags.DEFINE_string('master', 'local',
                    'BNS name of the TensorFlow master to use.')
flags.DEFINE_enum('accelerator_type', 'GPU', ['CPU', 'GPU', 'TPU'],
                  'Hardware type.')
flags.DEFINE_enum('dataset', 'imagenet', ['imagenet', 'custom'], 'Dataset')
flags.DEFINE_string('data_dir', None, 'Data directory for custom images.')
flags.DEFINE_string('tpu_worker_name', 'tpu_worker', 'Name of the TPU worker.')
flags.DEFINE_string(
    'pretrain_dir', None, 'Finetune from a pretrained checkpoint.')
flags.DEFINE_string('summaries_log_dir', 'summaries', 'Summaries parent.')
flags.DEFINE_integer('steps_per_summaries', 100, 'Steps per summaries.')
flags.DEFINE_integer('devices_per_worker', 1, 'Number of devices per worker.')
flags.DEFINE_integer('num_workers', 1, 'Number workers.')
config_flags.DEFINE_config_file(
    'config',
    default='test_configs/colorizer.py',
    help_string='Training configuration file.')

FLAGS = flags.FLAGS


def restore_checkpoint(model, ema, strategy, latest_ckpt=None, optimizer=None):
  if optimizer is None:
    ckpt_func = functools.partial(
        train_utils.create_checkpoint, models=model, ema=ema)
  else:
    ckpt_func = functools.partial(
        train_utils.create_checkpoint, models=model, ema=ema,
        optimizer=optimizer)

  checkpoint = train_utils.with_strategy(ckpt_func, strategy)
  if latest_ckpt:
    logging.info('Restoring from pretrained directory: %s', latest_ckpt)
    train_utils.with_strategy(lambda: checkpoint.restore(latest_ckpt), strategy)
  return checkpoint


def is_tpu():
  return FLAGS.accelerator_type == 'TPU'


def loss_on_batch(inputs, model, config, training=False):
  """Loss on a batch of inputs."""
  logits, aux_output = model.get_logits(
      inputs_dict=inputs, train_config=config, training=training)
  loss, aux_loss_dict = model.loss(
      targets=inputs, logits=logits, train_config=config, training=training,
      aux_output=aux_output)
  loss_factor = config.get('loss_factor', 1.0)

  loss_dict = collections.OrderedDict()
  loss_dict['loss'] = loss
  total_loss = loss_factor * loss

  for aux_key, aux_loss in aux_loss_dict.items():
    aux_loss_factor = config.get(f'{aux_key}_loss_factor', 1.0)
    loss_dict[aux_key] = aux_loss
    total_loss += aux_loss_factor * aux_loss
  loss_dict['total_loss'] = total_loss

  extra_info = collections.OrderedDict([
      ('scalar', loss_dict),
  ])
  return total_loss, extra_info


def train_step(config,
               model,
               optimizer,
               metrics,
               ema=None,
               strategy=None):
  """Training StepFn."""

  def step_fn(inputs):
    """Per-Replica StepFn."""
    with tf.GradientTape() as tape:
      loss, extra = loss_on_batch(inputs, model, config, training=True)
      scaled_loss = loss
      if strategy:
        scaled_loss /= float(strategy.num_replicas_in_sync)

    grads = tape.gradient(scaled_loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    for metric_key, metric in metrics.items():
      metric.update_state(extra['scalar'][metric_key])

    if ema is not None:
      ema.apply(model.trainable_variables)
    return loss

  return train_utils.step_with_strategy(step_fn, strategy)


def build(config, batch_size, is_train=False):
  optimizer = train_utils.build_optimizer(config)
  ema_vars = []

  downsample = config.get('downsample', False)
  downsample_res = config.get('downsample_res', 64)
  h, w = config.resolution

  if config.model.name == 'coltran_core':
    if downsample:
      h, w = downsample_res, downsample_res
    zero = tf.zeros((batch_size, h, w, 3), dtype=tf.int32)
    model = colorizer.ColTranCore(config.model)
    model(zero, training=is_train)

  c = 1 if is_train else 3
  if config.model.name == 'color_upsampler':
    if downsample:
      h, w = downsample_res, downsample_res
    zero_slice = tf.zeros((batch_size, h, w, c), dtype=tf.int32)
    zero = tf.zeros((batch_size, h, w, 3), dtype=tf.int32)
    model = upsampler.ColorUpsampler(config.model)
    model(zero, inputs_slice=zero_slice, training=is_train)
  elif config.model.name == 'spatial_upsampler':
    zero_slice = tf.zeros((batch_size, h, w, c), dtype=tf.int32)
    zero = tf.zeros((batch_size, h, w, 3), dtype=tf.int32)
    model = upsampler.SpatialUpsampler(config.model)
    model(zero, inputs_slice=zero_slice, training=is_train)

  ema_vars = model.trainable_variables
  ema = train_utils.build_ema(config, ema_vars)
  return model, optimizer, ema


###############################################################################
## Train.
###############################################################################
def train(logdir):
  config = FLAGS.config
  steps_per_write = FLAGS.steps_per_summaries
  train_utils.write_config(config, logdir)

  strategy, batch_size = train_utils.setup_strategy(
      config, FLAGS.master,
      FLAGS.devices_per_worker, FLAGS.mode, FLAGS.accelerator_type)

  def input_fn(input_context=None):
    read_config = None
    if input_context is not None:
      read_config = tfds.ReadConfig(input_context=input_context)

    dataset = datasets.get_dataset(
        name=FLAGS.dataset,
        config=config,
        batch_size=config.batch_size,
        subset='train',
        read_config=read_config,
        data_dir=FLAGS.data_dir)
    return dataset

  # DATASET CREATION.
  logging.info('Building dataset.')
  train_dataset = train_utils.dataset_with_strategy(input_fn, strategy)
  data_iterator = iter(train_dataset)

  # MODEL BUILDING
  logging.info('Building model.')
  model, optimizer, ema = train_utils.with_strategy(
      lambda: build(config, batch_size, True), strategy)
  model.summary(120, print_fn=logging.info)

  # METRIC CREATION.
  metrics = {}
  metric_keys = ['loss', 'total_loss']
  metric_keys += model.metric_keys
  for metric_key in metric_keys:
    func = functools.partial(tf.keras.metrics.Mean, metric_key)
    curr_metric = train_utils.with_strategy(func, strategy)
    metrics[metric_key] = curr_metric

  # CHECKPOINTING LOGIC.
  if FLAGS.pretrain_dir is not None:
    pretrain_ckpt = tf.train.latest_checkpoint(FLAGS.pretrain_dir)
    assert pretrain_ckpt

    # Load the entire model without the optimizer from the checkpoints.
    restore_checkpoint(model, ema, strategy, pretrain_ckpt, optimizer=None)
    # New tf.train.Checkpoint instance with a reset optimizer.
    checkpoint = restore_checkpoint(
        model, ema, strategy, latest_ckpt=None, optimizer=optimizer)
  else:
    latest_ckpt = tf.train.latest_checkpoint(logdir)
    checkpoint = restore_checkpoint(
        model, ema, strategy, latest_ckpt, optimizer=optimizer)

  checkpoint = tf.train.CheckpointManager(
      checkpoint, directory=logdir, checkpoint_name='model', max_to_keep=10)
  if optimizer.iterations.numpy() == 0:
    checkpoint_name = checkpoint.save()
    logging.info('Saved checkpoint to %s', checkpoint_name)

  train_summary_dir = os.path.join(logdir, 'train_summaries')
  writer = tf.summary.create_file_writer(train_summary_dir)
  start_time = time.time()

  logging.info('Start Training.')

  # This hack of wrapping up multiple train steps with a tf.function call
  # speeds up training significantly.
  # See: https://www.tensorflow.org/guide/tpu#improving_performance_by_multiple_steps_within_tffunction # pylint: disable=line-too-long
  @tf.function
  def train_multiple_steps(iterator, steps_per_epoch):

    train_step_f = train_step(config, model, optimizer, metrics, ema,
                              strategy)

    for _ in range(steps_per_epoch):
      train_step_f(iterator)

  while optimizer.iterations.numpy() < config.get('max_train_steps', 1000000):
    num_train_steps = optimizer.iterations

    for metric_key in metric_keys:
      metrics[metric_key].reset_states()

    start_run = time.time()

    train_multiple_steps(data_iterator, tf.convert_to_tensor(steps_per_write))

    steps_per_sec = steps_per_write / (time.time() - start_run)
    with writer.as_default():
      for metric_key, metric in metrics.items():
        metric_np = metric.result().numpy()
        tf.summary.scalar(metric_key, metric_np, step=num_train_steps)

        if metric_key == 'total_loss':
          logging.info('Loss: %.3f bits/dim, Speed: %.3f steps/second',
                       metric_np, steps_per_sec)

    if time.time() - start_time > config.save_checkpoint_secs:
      checkpoint_name = checkpoint.save()
      logging.info('Saved checkpoint to %s', checkpoint_name)
      start_time = time.time()


###############################################################################
## Evaluating.
###############################################################################


def evaluate(logdir, subset):
  """Executes the evaluation loop."""
  config = FLAGS.config
  strategy, batch_size = train_utils.setup_strategy(
      config, FLAGS.master,
      FLAGS.devices_per_worker, FLAGS.mode, FLAGS.accelerator_type)

  def input_fn(_=None):
    return datasets.get_dataset(
        name=config.dataset,
        config=config,
        batch_size=config.eval_batch_size,
        subset=subset)

  model, optimizer, ema = train_utils.with_strategy(
      lambda: build(config, batch_size, False), strategy)

  metric_keys = ['loss', 'total_loss']
  # metric_keys += model.metric_keys
  metrics = {}
  for metric_key in metric_keys:
    func = functools.partial(tf.keras.metrics.Mean, metric_key)
    curr_metric = train_utils.with_strategy(func, strategy)
    metrics[metric_key] = curr_metric

  checkpoints = train_utils.with_strategy(
      lambda: train_utils.create_checkpoint(model, optimizer, ema),
      strategy)
  dataset = train_utils.dataset_with_strategy(input_fn, strategy)

  def step_fn(batch):
    _, extra = loss_on_batch(batch, model, config, training=False)

    for metric_key in metric_keys:
      curr_metric = metrics[metric_key]
      curr_scalar = extra['scalar'][metric_key]
      curr_metric.update_state(curr_scalar)

  num_examples = config.eval_num_examples
  eval_step = train_utils.step_with_strategy(step_fn, strategy)
  ckpt_path = None
  wait_max = config.get(
      'eval_checkpoint_wait_secs', config.save_checkpoint_secs * 100)
  is_ema = True if ema else False

  eval_summary_dir = os.path.join(
      logdir, 'eval_{}_summaries_pyk_{}'.format(subset, is_ema))
  writer = tf.summary.create_file_writer(eval_summary_dir)

  while True:
    ckpt_path = train_utils.wait_for_checkpoint(logdir, ckpt_path, wait_max)
    logging.info(ckpt_path)
    if ckpt_path is None:
      logging.info('Timed out waiting for checkpoint.')
      break

    train_utils.with_strategy(
        lambda: train_utils.restore(model, checkpoints, logdir, ema),
        strategy)
    data_iterator = iter(dataset)
    num_steps = num_examples // batch_size

    for metric_key, metric in metrics.items():
      metric.reset_states()

    logging.info('Starting evaluation.')
    done = False
    for i in range(0, num_steps, FLAGS.steps_per_summaries):
      start_run = time.time()
      for k in range(min(num_steps - i, FLAGS.steps_per_summaries)):
        try:
          if k % 10 == 0:
            logging.info('Step: %d', (i + k + 1))
          eval_step(data_iterator)
        except (StopIteration, tf.errors.OutOfRangeError):
          done = True
          break
      if done:
        break
      bits_per_dim = metrics['loss'].result()
      logging.info('Bits/Dim: %.3f, Speed: %.3f seconds/step, Step: %d/%d',
                   bits_per_dim,
                   (time.time() - start_run) / FLAGS.steps_per_summaries,
                   i + k + 1, num_steps)

    # logging.info('Final Bits/Dim: %.3f', bits_per_dim)
    with writer.as_default():
      for metric_key, metric in metrics.items():
        curr_scalar = metric.result().numpy()
        tf.summary.scalar(metric_key, curr_scalar, step=optimizer.iterations)


def main(_):
  logging.info('Logging to %s.', FLAGS.logdir)
  if FLAGS.mode == 'train':
    logging.info('[main] I am the trainer.')
    try:
      train(FLAGS.logdir)
    # During TPU Preemeption, the coordinator hangs with the error below.
    # the exception forces the coordinator to fail, and it will be restarted.
    except (tf.errors.UnavailableError, tf.errors.CancelledError):
      os._exit(os.EX_TEMPFAIL)  # pylint: disable=protected-access
  elif FLAGS.mode.startswith('train'):
    logging.info('[main] I am the trainer.')
    train(os.path.join(FLAGS.logdir, FLAGS.mode))
  elif FLAGS.mode == 'eval_train':
    logging.info('[main] I am the training set evaluator.')
    evaluate(FLAGS.logdir, subset='train')
  elif FLAGS.mode == 'eval_valid':
    logging.info('[main] I am the validation set evaluator.')
    evaluate(FLAGS.logdir, subset='valid')
  elif FLAGS.mode == 'eval_test':
    logging.info('[main] I am the test set evaluator.')
    evaluate(FLAGS.logdir, subset='test')
  else:
    raise ValueError(
        'Unknown mode {}. '
        'Must be one of [train, eval_train, eval_valid, eval_test]'.format(
            FLAGS.mode))


if __name__ == '__main__':
  app.run(main)

[(47, 'absl.flags.DEFINE_enum', 'flags.DEFINE_enum', (['"""mode"""', '"""train"""', "['train', 'eval_train', 'eval_valid', 'eval_test']", '"""Operation mode."""'], {}), False, 'from absl import flags\n'), (50, 'absl.flags.DEFINE_string', 'flags.DEFINE_string', (['"""logdir"""', '"""/tmp/svt"""', '"""Main directory for logs."""'], {}), False, 'from absl import flags\n'), (51, 'absl.flags.DEFINE_string', 'flags.DEFINE_string', (['"""master"""', '"""local"""', '"""BNS name of the TensorFlow master to use."""'], {}), False, 'from absl import flags\n'), (53, 'absl.flags.DEFINE_enum', 'flags.DEFINE_enum', (['"""accelerator_type"""', '"""GPU"""', "['CPU', 'GPU', 'TPU']", '"""Hardware type."""'], {}), False, 'from absl import flags\n'), (55, 'absl.flags.DEFINE_enum', 'flags.DEFINE_enum', (['"""dataset"""', '"""imagenet"""', "['imagenet', 'custom']", '"""Dataset"""'], {}), False, 'from absl import flags\n'), (56, 'absl.flags.DEFINE_string', 'flags.DEFINE_string', (['"""data_dir"""', 'None', '"""Data directory for custom images."""'], {}), False, 'from absl import flags\n'), (57, 'absl.flags.DEFINE_string', 'flags.DEFINE_string', (['"""tpu_worker_name"""', '"""tpu_worker"""', '"""Name of the TPU worker."""'], {}), False, 'from absl import flags\n'), (58, 'absl.flags.DEFINE_string', 'flags.DEFINE_string', (['"""pretrain_dir"""', 'None', '"""Finetune from a pretrained checkpoint."""'], {}), False, 'from absl import flags\n'), (60, 'absl.flags.DEFINE_string', 'flags.DEFINE_string', (['"""summaries_log_dir"""', '"""summaries"""', '"""Summaries parent."""'], {}), False, 'from absl import flags\n'), (61, 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', (['"""steps_per_summaries"""', '(100)', '"""Steps per summaries."""'], {}), False, 'from absl import flags\n'), (62, 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', (['"""devices_per_worker"""', '(1)', '"""Number of devices per worker."""'], {}), False, 'from absl import flags\n'), (63, 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', (['"""num_workers"""', '(1)', '"""Number workers."""'], {}), False, 'from absl import flags\n'), (64, 'ml_collections.config_flags.DEFINE_config_file', 'config_flags.DEFINE_config_file', (['"""config"""'], {'default': '"""test_configs/colorizer.py"""', 'help_string': '"""Training configuration file."""'}), False, 'from ml_collections import config_flags\n'), (81, 'coltran.utils.train_utils.with_strategy', 'train_utils.with_strategy', (['ckpt_func', 'strategy'], {}), False, 'from coltran.utils import train_utils\n'), (101, 'collections.OrderedDict', 'collections.OrderedDict', ([], {}), False, 'import collections\n'), (111, 'collections.OrderedDict', 'collections.OrderedDict', (["[('scalar', loss_dict)]"], {}), False, 'import collections\n'), (142, 'coltran.utils.train_utils.step_with_strategy', 'train_utils.step_with_strategy', (['step_fn', 'strategy'], {}), False, 'from coltran.utils import train_utils\n'), (146, 'coltran.utils.train_utils.build_optimizer', 'train_utils.build_optimizer', (['config'], {}), False, 'from coltran.utils import train_utils\n'), (175, 'coltran.utils.train_utils.build_ema', 'train_utils.build_ema', (['config', 'ema_vars'], {}), False, 'from coltran.utils import train_utils\n'), (185, 'coltran.utils.train_utils.write_config', 'train_utils.write_config', (['config', 'logdir'], {}), False, 'from coltran.utils import train_utils\n'), (187, 'coltran.utils.train_utils.setup_strategy', 'train_utils.setup_strategy', (['config', 'FLAGS.master', 'FLAGS.devices_per_worker', 'FLAGS.mode', 'FLAGS.accelerator_type'], {}), False, 'from coltran.utils import train_utils\n'), (206, 'absl.logging.info', 'logging.info', (['"""Building dataset."""'], {}), False, 'from absl import logging\n'), (207, 'coltran.utils.train_utils.dataset_with_strategy', 'train_utils.dataset_with_strategy', (['input_fn', 'strategy'], {}), False, 'from coltran.utils import train_utils\n'), (211, 'absl.logging.info', 'logging.info', (['"""Building model."""'], {}), False, 'from absl import logging\n'), (240, 'tensorflow.train.CheckpointManager', 'tf.train.CheckpointManager', (['checkpoint'], {'directory': 'logdir', 'checkpoint_name': '"""model"""', 'max_to_keep': '(10)'}), True, 'import tensorflow as tf\n'), (246, 'os.path.join', 'os.path.join', (['logdir', '"""train_summaries"""'], {}), False, 'import os\n'), (247, 'tensorflow.summary.create_file_writer', 'tf.summary.create_file_writer', (['train_summary_dir'], {}), True, 'import tensorflow as tf\n'), (248, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (250, 'absl.logging.info', 'logging.info', (['"""Start Training."""'], {}), False, 'from absl import logging\n'), (298, 'coltran.utils.train_utils.setup_strategy', 'train_utils.setup_strategy', (['config', 'FLAGS.master', 'FLAGS.devices_per_worker', 'FLAGS.mode', 'FLAGS.accelerator_type'], {}), False, 'from coltran.utils import train_utils\n'), (323, 'coltran.utils.train_utils.dataset_with_strategy', 'train_utils.dataset_with_strategy', (['input_fn', 'strategy'], {}), False, 'from coltran.utils import train_utils\n'), (334, 'coltran.utils.train_utils.step_with_strategy', 'train_utils.step_with_strategy', (['step_fn', 'strategy'], {}), False, 'from coltran.utils import train_utils\n'), (342, 'tensorflow.summary.create_file_writer', 'tf.summary.create_file_writer', (['eval_summary_dir'], {}), True, 'import tensorflow as tf\n'), (388, 'absl.logging.info', 'logging.info', (['"""Logging to %s."""', 'FLAGS.logdir'], {}), False, 'from absl import logging\n'), (417, 'absl.app.run', 'app.run', (['main'], {}), False, 'from absl import app\n'), (74, 'functools.partial', 'functools.partial', (['train_utils.create_checkpoint'], {'models': 'model', 'ema': 'ema'}), False, 'import functools\n'), (77, 'functools.partial', 'functools.partial', (['train_utils.create_checkpoint'], {'models': 'model', 'ema': 'ema', 'optimizer': 'optimizer'}), False, 'import functools\n'), (83, 'absl.logging.info', 'logging.info', (['"""Restoring from pretrained directory: %s"""', 'latest_ckpt'], {}), False, 'from absl import logging\n'), (156, 'tensorflow.zeros', 'tf.zeros', (['(batch_size, h, w, 3)'], {'dtype': 'tf.int32'}), True, 'import tensorflow as tf\n'), (157, 'coltran.models.colorizer.ColTranCore', 'colorizer.ColTranCore', (['config.model'], {}), False, 'from coltran.models import colorizer\n'), (164, 'tensorflow.zeros', 'tf.zeros', (['(batch_size, h, w, c)'], {'dtype': 'tf.int32'}), True, 'import tensorflow as tf\n'), (165, 'tensorflow.zeros', 'tf.zeros', (['(batch_size, h, w, 3)'], {'dtype': 'tf.int32'}), True, 'import tensorflow as tf\n'), (166, 'coltran.models.upsampler.ColorUpsampler', 'upsampler.ColorUpsampler', (['config.model'], {}), False, 'from coltran.models import upsampler\n'), (196, 'coltran.datasets.get_dataset', 'datasets.get_dataset', ([], {'name': 'FLAGS.dataset', 'config': 'config', 'batch_size': 'config.batch_size', 'subset': '"""train"""', 'read_config': 'read_config', 'data_dir': 'FLAGS.data_dir'}), False, 'from coltran import datasets\n'), (221, 'functools.partial', 'functools.partial', (['tf.keras.metrics.Mean', 'metric_key'], {}), False, 'import functools\n'), (222, 'coltran.utils.train_utils.with_strategy', 'train_utils.with_strategy', (['func', 'strategy'], {}), False, 'from coltran.utils import train_utils\n'), (227, 'tensorflow.train.latest_checkpoint', 'tf.train.latest_checkpoint', (['FLAGS.pretrain_dir'], {}), True, 'import tensorflow as tf\n'), (236, 'tensorflow.train.latest_checkpoint', 'tf.train.latest_checkpoint', (['logdir'], {}), True, 'import tensorflow as tf\n'), (244, 'absl.logging.info', 'logging.info', (['"""Saved checkpoint to %s"""', 'checkpoint_name'], {}), False, 'from absl import logging\n'), (270, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (303, 'coltran.datasets.get_dataset', 'datasets.get_dataset', ([], {'name': 'config.dataset', 'config': 'config', 'batch_size': 'config.eval_batch_size', 'subset': 'subset'}), False, 'from coltran import datasets\n'), (316, 'functools.partial', 'functools.partial', (['tf.keras.metrics.Mean', 'metric_key'], {}), False, 'import functools\n'), (317, 'coltran.utils.train_utils.with_strategy', 'train_utils.with_strategy', (['func', 'strategy'], {}), False, 'from coltran.utils import train_utils\n'), (345, 'coltran.utils.train_utils.wait_for_checkpoint', 'train_utils.wait_for_checkpoint', (['logdir', 'ckpt_path', 'wait_max'], {}), False, 'from coltran.utils import train_utils\n'), (346, 'absl.logging.info', 'logging.info', (['ckpt_path'], {}), False, 'from absl import logging\n'), (360, 'absl.logging.info', 'logging.info', (['"""Starting evaluation."""'], {}), False, 'from absl import logging\n'), (390, 'absl.logging.info', 'logging.info', (['"""[main] I am the trainer."""'], {}), False, 'from absl import logging\n'), (127, 'tensorflow.GradientTape', 'tf.GradientTape', ([], {}), True, 'import tensorflow as tf\n'), (169, 'tensorflow.zeros', 'tf.zeros', (['(batch_size, h, w, c)'], {'dtype': 'tf.int32'}), True, 'import tensorflow as tf\n'), (170, 'tensorflow.zeros', 'tf.zeros', (['(batch_size, h, w, 3)'], {'dtype': 'tf.int32'}), True, 'import tensorflow as tf\n'), (171, 'coltran.models.upsampler.SpatialUpsampler', 'upsampler.SpatialUpsampler', (['config.model'], {}), False, 'from coltran.models import upsampler\n'), (194, 'tensorflow_datasets.ReadConfig', 'tfds.ReadConfig', ([], {'input_context': 'input_context'}), True, 'import tensorflow_datasets as tfds\n'), (272, 'tensorflow.convert_to_tensor', 'tf.convert_to_tensor', (['steps_per_write'], {}), True, 'import tensorflow as tf\n'), (286, 'absl.logging.info', 'logging.info', (['"""Saved checkpoint to %s"""', 'checkpoint_name'], {}), False, 'from absl import logging\n'), (287, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (321, 'coltran.utils.train_utils.create_checkpoint', 'train_utils.create_checkpoint', (['model', 'optimizer', 'ema'], {}), False, 'from coltran.utils import train_utils\n'), (348, 'absl.logging.info', 'logging.info', (['"""Timed out waiting for checkpoint."""'], {}), False, 'from absl import logging\n'), (363, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (398, 'absl.logging.info', 'logging.info', (['"""[main] I am the trainer."""'], {}), False, 'from absl import logging\n'), (274, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (278, 'tensorflow.summary.scalar', 'tf.summary.scalar', (['metric_key', 'metric_np'], {'step': 'num_train_steps'}), True, 'import tensorflow as tf\n'), (284, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (352, 'coltran.utils.train_utils.restore', 'train_utils.restore', (['model', 'checkpoints', 'logdir', 'ema'], {}), False, 'from coltran.utils import train_utils\n'), (384, 'tensorflow.summary.scalar', 'tf.summary.scalar', (['metric_key', 'curr_scalar'], {'step': 'optimizer.iterations'}), True, 'import tensorflow as tf\n'), (396, 'os._exit', 'os._exit', (['os.EX_TEMPFAIL'], {}), False, 'import os\n'), (399, 'os.path.join', 'os.path.join', (['FLAGS.logdir', 'FLAGS.mode'], {}), False, 'import os\n'), (401, 'absl.logging.info', 'logging.info', (['"""[main] I am the training set evaluator."""'], {}), False, 'from absl import logging\n'), (281, 'absl.logging.info', 'logging.info', (['"""Loss: %.3f bits/dim, Speed: %.3f steps/second"""', 'metric_np', 'steps_per_sec'], {}), False, 'from absl import logging\n'), (404, 'absl.logging.info', 'logging.info', (['"""[main] I am the validation set evaluator."""'], {}), False, 'from absl import logging\n'), (367, 'absl.logging.info', 'logging.info', (['"""Step: %d"""', '(i + k + 1)'], {}), False, 'from absl import logging\n'), (377, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (407, 'absl.logging.info', 'logging.info', (['"""[main] I am the test set evaluator."""'], {}), False, 'from absl import logging\n')]