# 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)

[((1288, 1397), 'absl.flags.DEFINE_enum', 'flags.DEFINE_enum', (['"""mode"""', '"""train"""', "['train', 'eval_train', 'eval_valid', 'eval_test']", '"""Operation mode."""'], {}), "('mode', 'train', ['train', 'eval_train', 'eval_valid',\n    'eval_test'], 'Operation mode.')\n", (1305, 1397), False, 'from absl import flags\n'), ((1400, 1469), 'absl.flags.DEFINE_string', 'flags.DEFINE_string', (['"""logdir"""', '"""/tmp/svt"""', '"""Main directory for logs."""'], {}), "('logdir', '/tmp/svt', 'Main directory for logs.')\n", (1419, 1469), False, 'from absl import flags\n'), ((1470, 1557), 'absl.flags.DEFINE_string', 'flags.DEFINE_string', (['"""master"""', '"""local"""', '"""BNS name of the TensorFlow master to use."""'], {}), "('master', 'local',\n    'BNS name of the TensorFlow master to use.')\n", (1489, 1557), False, 'from absl import flags\n'), ((1574, 1663), 'absl.flags.DEFINE_enum', 'flags.DEFINE_enum', (['"""accelerator_type"""', '"""GPU"""', "['CPU', 'GPU', 'TPU']", '"""Hardware type."""'], {}), "('accelerator_type', 'GPU', ['CPU', 'GPU', 'TPU'],\n    'Hardware type.')\n", (1591, 1663), False, 'from absl import flags\n'), ((1678, 1753), 'absl.flags.DEFINE_enum', 'flags.DEFINE_enum', (['"""dataset"""', '"""imagenet"""', "['imagenet', 'custom']", '"""Dataset"""'], {}), "('dataset', 'imagenet', ['imagenet', 'custom'], 'Dataset')\n", (1695, 1753), False, 'from absl import flags\n'), ((1754, 1828), 'absl.flags.DEFINE_string', 'flags.DEFINE_string', (['"""data_dir"""', 'None', '"""Data directory for custom images."""'], {}), "('data_dir', None, 'Data directory for custom images.')\n", (1773, 1828), False, 'from absl import flags\n'), ((1829, 1908), 'absl.flags.DEFINE_string', 'flags.DEFINE_string', (['"""tpu_worker_name"""', '"""tpu_worker"""', '"""Name of the TPU worker."""'], {}), "('tpu_worker_name', 'tpu_worker', 'Name of the TPU worker.')\n", (1848, 1908), False, 'from absl import flags\n'), ((1909, 1996), 'absl.flags.DEFINE_string', 'flags.DEFINE_string', (['"""pretrain_dir"""', 'None', '"""Finetune from a pretrained checkpoint."""'], {}), "('pretrain_dir', None,\n    'Finetune from a pretrained checkpoint.')\n", (1928, 1996), False, 'from absl import flags\n'), ((1998, 2072), 'absl.flags.DEFINE_string', 'flags.DEFINE_string', (['"""summaries_log_dir"""', '"""summaries"""', '"""Summaries parent."""'], {}), "('summaries_log_dir', 'summaries', 'Summaries parent.')\n", (2017, 2072), False, 'from absl import flags\n'), ((2073, 2145), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', (['"""steps_per_summaries"""', '(100)', '"""Steps per summaries."""'], {}), "('steps_per_summaries', 100, 'Steps per summaries.')\n", (2093, 2145), False, 'from absl import flags\n'), ((2146, 2224), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', (['"""devices_per_worker"""', '(1)', '"""Number of devices per worker."""'], {}), "('devices_per_worker', 1, 'Number of devices per worker.')\n", (2166, 2224), False, 'from absl import flags\n'), ((2225, 2282), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', (['"""num_workers"""', '(1)', '"""Number workers."""'], {}), "('num_workers', 1, 'Number workers.')\n", (2245, 2282), False, 'from absl import flags\n'), ((2283, 2410), 'ml_collections.config_flags.DEFINE_config_file', 'config_flags.DEFINE_config_file', (['"""config"""'], {'default': '"""test_configs/colorizer.py"""', 'help_string': '"""Training configuration file."""'}), "('config', default=\n    'test_configs/colorizer.py', help_string='Training configuration file.')\n", (2314, 2410), False, 'from ml_collections import config_flags\n'), ((2793, 2839), 'coltran.utils.train_utils.with_strategy', 'train_utils.with_strategy', (['ckpt_func', 'strategy'], {}), '(ckpt_func, strategy)\n', (2818, 2839), False, 'from coltran.utils import train_utils\n'), ((3495, 3520), 'collections.OrderedDict', 'collections.OrderedDict', ([], {}), '()\n', (3518, 3520), False, 'import collections\n'), ((3831, 3879), 'collections.OrderedDict', 'collections.OrderedDict', (["[('scalar', loss_dict)]"], {}), "([('scalar', loss_dict)])\n", (3854, 3879), False, 'import collections\n'), ((4699, 4748), 'coltran.utils.train_utils.step_with_strategy', 'train_utils.step_with_strategy', (['step_fn', 'strategy'], {}), '(step_fn, strategy)\n', (4729, 4748), False, 'from coltran.utils import train_utils\n'), ((4812, 4847), 'coltran.utils.train_utils.build_optimizer', 'train_utils.build_optimizer', (['config'], {}), '(config)\n', (4839, 4847), False, 'from coltran.utils import train_utils\n'), ((5944, 5983), 'coltran.utils.train_utils.build_ema', 'train_utils.build_ema', (['config', 'ema_vars'], {}), '(config, ema_vars)\n', (5965, 5983), False, 'from coltran.utils import train_utils\n'), ((6278, 6318), 'coltran.utils.train_utils.write_config', 'train_utils.write_config', (['config', 'logdir'], {}), '(config, logdir)\n', (6302, 6318), False, 'from coltran.utils import train_utils\n'), ((6345, 6459), 'coltran.utils.train_utils.setup_strategy', 'train_utils.setup_strategy', (['config', 'FLAGS.master', 'FLAGS.devices_per_worker', 'FLAGS.mode', 'FLAGS.accelerator_type'], {}), '(config, FLAGS.master, FLAGS.devices_per_worker,\n    FLAGS.mode, FLAGS.accelerator_type)\n', (6371, 6459), False, 'from coltran.utils import train_utils\n'), ((6888, 6921), 'absl.logging.info', 'logging.info', (['"""Building dataset."""'], {}), "('Building dataset.')\n", (6900, 6921), False, 'from absl import logging\n'), ((6940, 6993), 'coltran.utils.train_utils.dataset_with_strategy', 'train_utils.dataset_with_strategy', (['input_fn', 'strategy'], {}), '(input_fn, strategy)\n', (6973, 6993), False, 'from coltran.utils import train_utils\n'), ((7054, 7085), 'absl.logging.info', 'logging.info', (['"""Building model."""'], {}), "('Building model.')\n", (7066, 7085), False, 'from absl import logging\n'), ((8197, 8299), 'tensorflow.train.CheckpointManager', 'tf.train.CheckpointManager', (['checkpoint'], {'directory': 'logdir', 'checkpoint_name': '"""model"""', 'max_to_keep': '(10)'}), "(checkpoint, directory=logdir, checkpoint_name=\n    'model', max_to_keep=10)\n", (8223, 8299), True, 'import tensorflow as tf\n'), ((8465, 8504), 'os.path.join', 'os.path.join', (['logdir', '"""train_summaries"""'], {}), "(logdir, 'train_summaries')\n", (8477, 8504), False, 'import os\n'), ((8516, 8564), 'tensorflow.summary.create_file_writer', 'tf.summary.create_file_writer', (['train_summary_dir'], {}), '(train_summary_dir)\n', (8545, 8564), True, 'import tensorflow as tf\n'), ((8580, 8591), 'time.time', 'time.time', ([], {}), '()\n', (8589, 8591), False, 'import time\n'), ((8595, 8626), 'absl.logging.info', 'logging.info', (['"""Start Training."""'], {}), "('Start Training.')\n", (8607, 8626), False, 'from absl import logging\n'), ((10348, 10462), 'coltran.utils.train_utils.setup_strategy', 'train_utils.setup_strategy', (['config', 'FLAGS.master', 'FLAGS.devices_per_worker', 'FLAGS.mode', 'FLAGS.accelerator_type'], {}), '(config, FLAGS.master, FLAGS.devices_per_worker,\n    FLAGS.mode, FLAGS.accelerator_type)\n', (10374, 10462), False, 'from coltran.utils import train_utils\n'), ((11187, 11240), 'coltran.utils.train_utils.dataset_with_strategy', 'train_utils.dataset_with_strategy', (['input_fn', 'strategy'], {}), '(input_fn, strategy)\n', (11220, 11240), False, 'from coltran.utils import train_utils\n'), ((11556, 11605), 'coltran.utils.train_utils.step_with_strategy', 'train_utils.step_with_strategy', (['step_fn', 'strategy'], {}), '(step_fn, strategy)\n', (11586, 11605), False, 'from coltran.utils import train_utils\n'), ((11866, 11913), 'tensorflow.summary.create_file_writer', 'tf.summary.create_file_writer', (['eval_summary_dir'], {}), '(eval_summary_dir)\n', (11895, 11913), True, 'import tensorflow as tf\n'), ((13418, 13462), 'absl.logging.info', 'logging.info', (['"""Logging to %s."""', 'FLAGS.logdir'], {}), "('Logging to %s.', FLAGS.logdir)\n", (13430, 13462), False, 'from absl import logging\n'), ((14591, 14604), 'absl.app.run', 'app.run', (['main'], {}), '(main)\n', (14598, 14604), False, 'from absl import app\n'), ((2562, 2633), 'functools.partial', 'functools.partial', (['train_utils.create_checkpoint'], {'models': 'model', 'ema': 'ema'}), '(train_utils.create_checkpoint, models=model, ema=ema)\n', (2579, 2633), False, 'import functools\n'), ((2667, 2763), 'functools.partial', 'functools.partial', (['train_utils.create_checkpoint'], {'models': 'model', 'ema': 'ema', 'optimizer': 'optimizer'}), '(train_utils.create_checkpoint, models=model, ema=ema,\n    optimizer=optimizer)\n', (2684, 2763), False, 'import functools\n'), ((2862, 2930), 'absl.logging.info', 'logging.info', (['"""Restoring from pretrained directory: %s"""', 'latest_ckpt'], {}), "('Restoring from pretrained directory: %s', latest_ckpt)\n", (2874, 2930), False, 'from absl import logging\n'), ((5108, 5155), 'tensorflow.zeros', 'tf.zeros', (['(batch_size, h, w, 3)'], {'dtype': 'tf.int32'}), '((batch_size, h, w, 3), dtype=tf.int32)\n', (5116, 5155), True, 'import tensorflow as tf\n'), ((5168, 5203), 'coltran.models.colorizer.ColTranCore', 'colorizer.ColTranCore', (['config.model'], {}), '(config.model)\n', (5189, 5203), False, 'from coltran.models import colorizer\n'), ((5392, 5439), 'tensorflow.zeros', 'tf.zeros', (['(batch_size, h, w, c)'], {'dtype': 'tf.int32'}), '((batch_size, h, w, c), dtype=tf.int32)\n', (5400, 5439), True, 'import tensorflow as tf\n'), ((5451, 5498), 'tensorflow.zeros', 'tf.zeros', (['(batch_size, h, w, 3)'], {'dtype': 'tf.int32'}), '((batch_size, h, w, 3), dtype=tf.int32)\n', (5459, 5498), True, 'import tensorflow as tf\n'), ((5511, 5549), 'coltran.models.upsampler.ColorUpsampler', 'upsampler.ColorUpsampler', (['config.model'], {}), '(config.model)\n', (5535, 5549), False, 'from coltran.models import upsampler\n'), ((6643, 6804), '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'}), "(name=FLAGS.dataset, config=config, batch_size=config.\n    batch_size, subset='train', read_config=read_config, data_dir=FLAGS.\n    data_dir)\n", (6663, 6804), False, 'from coltran import datasets\n'), ((7395, 7447), 'functools.partial', 'functools.partial', (['tf.keras.metrics.Mean', 'metric_key'], {}), '(tf.keras.metrics.Mean, metric_key)\n', (7412, 7447), False, 'import functools\n'), ((7466, 7507), 'coltran.utils.train_utils.with_strategy', 'train_utils.with_strategy', (['func', 'strategy'], {}), '(func, strategy)\n', (7491, 7507), False, 'from coltran.utils import train_utils\n'), ((7629, 7675), 'tensorflow.train.latest_checkpoint', 'tf.train.latest_checkpoint', (['FLAGS.pretrain_dir'], {}), '(FLAGS.pretrain_dir)\n', (7655, 7675), True, 'import tensorflow as tf\n'), ((8045, 8079), 'tensorflow.train.latest_checkpoint', 'tf.train.latest_checkpoint', (['logdir'], {}), '(logdir)\n', (8071, 8079), True, 'import tensorflow as tf\n'), ((8386, 8441), 'absl.logging.info', 'logging.info', (['"""Saved checkpoint to %s"""', 'checkpoint_name'], {}), "('Saved checkpoint to %s', checkpoint_name)\n", (8398, 8441), False, 'from absl import logging\n'), ((9341, 9352), 'time.time', 'time.time', ([], {}), '()\n', (9350, 9352), False, 'import time\n'), ((10508, 10619), 'coltran.datasets.get_dataset', 'datasets.get_dataset', ([], {'name': 'config.dataset', 'config': 'config', 'batch_size': 'config.eval_batch_size', 'subset': 'subset'}), '(name=config.dataset, config=config, batch_size=config.\n    eval_batch_size, subset=subset)\n', (10528, 10619), False, 'from coltran import datasets\n'), ((10896, 10948), 'functools.partial', 'functools.partial', (['tf.keras.metrics.Mean', 'metric_key'], {}), '(tf.keras.metrics.Mean, metric_key)\n', (10913, 10948), False, 'import functools\n'), ((10967, 11008), 'coltran.utils.train_utils.with_strategy', 'train_utils.with_strategy', (['func', 'strategy'], {}), '(func, strategy)\n', (10992, 11008), False, 'from coltran.utils import train_utils\n'), ((11945, 12005), 'coltran.utils.train_utils.wait_for_checkpoint', 'train_utils.wait_for_checkpoint', (['logdir', 'ckpt_path', 'wait_max'], {}), '(logdir, ckpt_path, wait_max)\n', (11976, 12005), False, 'from coltran.utils import train_utils\n'), ((12010, 12033), 'absl.logging.info', 'logging.info', (['ckpt_path'], {}), '(ckpt_path)\n', (12022, 12033), False, 'from absl import logging\n'), ((12406, 12442), 'absl.logging.info', 'logging.info', (['"""Starting evaluation."""'], {}), "('Starting evaluation.')\n", (12418, 12442), False, 'from absl import logging\n'), ((13495, 13535), 'absl.logging.info', 'logging.info', (['"""[main] I am the trainer."""'], {}), "('[main] I am the trainer.')\n", (13507, 13535), False, 'from absl import logging\n'), ((4164, 4181), 'tensorflow.GradientTape', 'tf.GradientTape', ([], {}), '()\n', (4179, 4181), True, 'import tensorflow as tf\n'), ((5676, 5723), 'tensorflow.zeros', 'tf.zeros', (['(batch_size, h, w, c)'], {'dtype': 'tf.int32'}), '((batch_size, h, w, c), dtype=tf.int32)\n', (5684, 5723), True, 'import tensorflow as tf\n'), ((5735, 5782), 'tensorflow.zeros', 'tf.zeros', (['(batch_size, h, w, 3)'], {'dtype': 'tf.int32'}), '((batch_size, h, w, 3), dtype=tf.int32)\n', (5743, 5782), True, 'import tensorflow as tf\n'), ((5795, 5835), 'coltran.models.upsampler.SpatialUpsampler', 'upsampler.SpatialUpsampler', (['config.model'], {}), '(config.model)\n', (5821, 5835), False, 'from coltran.models import upsampler\n'), ((6583, 6627), 'tensorflow_datasets.ReadConfig', 'tfds.ReadConfig', ([], {'input_context': 'input_context'}), '(input_context=input_context)\n', (6598, 6627), True, 'import tensorflow_datasets as tfds\n'), ((9394, 9431), 'tensorflow.convert_to_tensor', 'tf.convert_to_tensor', (['steps_per_write'], {}), '(steps_per_write)\n', (9414, 9431), True, 'import tensorflow as tf\n'), ((9965, 10020), 'absl.logging.info', 'logging.info', (['"""Saved checkpoint to %s"""', 'checkpoint_name'], {}), "('Saved checkpoint to %s', checkpoint_name)\n", (9977, 10020), False, 'from absl import logging\n'), ((10040, 10051), 'time.time', 'time.time', ([], {}), '()\n', (10049, 10051), False, 'import time\n'), ((11105, 11157), 'coltran.utils.train_utils.create_checkpoint', 'train_utils.create_checkpoint', (['model', 'optimizer', 'ema'], {}), '(model, optimizer, ema)\n', (11134, 11157), False, 'from coltran.utils import train_utils\n'), ((12066, 12115), 'absl.logging.info', 'logging.info', (['"""Timed out waiting for checkpoint."""'], {}), "('Timed out waiting for checkpoint.')\n", (12078, 12115), False, 'from absl import logging\n'), ((12539, 12550), 'time.time', 'time.time', ([], {}), '()\n', (12548, 12550), False, 'import time\n'), ((13900, 13940), 'absl.logging.info', 'logging.info', (['"""[main] I am the trainer."""'], {}), "('[main] I am the trainer.')\n", (13912, 13940), False, 'from absl import logging\n'), ((9473, 9484), 'time.time', 'time.time', ([], {}), '()\n', (9482, 9484), False, 'import time\n'), ((9629, 9691), 'tensorflow.summary.scalar', 'tf.summary.scalar', (['metric_key', 'metric_np'], {'step': 'num_train_steps'}), '(metric_key, metric_np, step=num_train_steps)\n', (9646, 9691), True, 'import tensorflow as tf\n'), ((9861, 9872), 'time.time', 'time.time', ([], {}), '()\n', (9870, 9872), False, 'import time\n'), ((12176, 12228), 'coltran.utils.train_utils.restore', 'train_utils.restore', (['model', 'checkpoints', 'logdir', 'ema'], {}), '(model, checkpoints, logdir, ema)\n', (12195, 12228), False, 'from coltran.utils import train_utils\n'), ((13331, 13400), 'tensorflow.summary.scalar', 'tf.summary.scalar', (['metric_key', 'curr_scalar'], {'step': 'optimizer.iterations'}), '(metric_key, curr_scalar, step=optimizer.iterations)\n', (13348, 13400), True, 'import tensorflow as tf\n'), ((13796, 13820), 'os._exit', 'os._exit', (['os.EX_TEMPFAIL'], {}), '(os.EX_TEMPFAIL)\n', (13804, 13820), False, 'import os\n'), ((13951, 13989), 'os.path.join', 'os.path.join', (['FLAGS.logdir', 'FLAGS.mode'], {}), '(FLAGS.logdir, FLAGS.mode)\n', (13963, 13989), False, 'import os\n'), ((14030, 14085), 'absl.logging.info', 'logging.info', (['"""[main] I am the training set evaluator."""'], {}), "('[main] I am the training set evaluator.')\n", (14042, 14085), False, 'from absl import logging\n'), ((9742, 9833), 'absl.logging.info', 'logging.info', (['"""Loss: %.3f bits/dim, Speed: %.3f steps/second"""', 'metric_np', 'steps_per_sec'], {}), "('Loss: %.3f bits/dim, Speed: %.3f steps/second', metric_np,\n    steps_per_sec)\n", (9754, 9833), False, 'from absl import logging\n'), ((14168, 14225), 'absl.logging.info', 'logging.info', (['"""[main] I am the validation set evaluator."""'], {}), "('[main] I am the validation set evaluator.')\n", (14180, 14225), False, 'from absl import logging\n'), ((12671, 12706), 'absl.logging.info', 'logging.info', (['"""Step: %d"""', '(i + k + 1)'], {}), "('Step: %d', i + k + 1)\n", (12683, 12706), False, 'from absl import logging\n'), ((13045, 13056), 'time.time', 'time.time', ([], {}), '()\n', (13054, 13056), False, 'import time\n'), ((14307, 14358), 'absl.logging.info', 'logging.info', (['"""[main] I am the test set evaluator."""'], {}), "('[main] I am the test set evaluator.')\n", (14319, 14358), False, 'from absl import logging\n')]