def train(*, rng: jnp.ndarray, config: ml_collections.ConfigDict,
          dataset: dataset_utils.Dataset, workdir: str,
          writer: metric_writers.MetricWriter):
  """Main training loop lives in this function."""
  lead_host = jax.process_index() == 0
  dtype = train_utils.get_dtype(config)
  writer.write_hparams(eval_utils.flatten_config(config))

  # Build the flax_model and the optimizers.
  _, init_rng = jax.random.split(rng)
  input_spec = [(
      dataset.meta_data['input_shape'],  # bs, t, h, w, 3
      dataset.meta_data.get('input_dtype', dtype))]
  model_dict, train_state, learning_rate_fn = create_train_state(
      input_spec, config, init_rng, True)

  # Restores the model if needed.
  if config.logging.enable_checkpoint:
    train_state = checkpoints.restore_checkpoint(
        ckpt_dir=workdir, target=train_state)

  start_step = int(train_state.global_step)
  # Log model size
  if start_step == 0:
    if config.get('vqgan.finetune_path'):
      logging.info(
          'loading the model for finetuning from %s', config.vqgan.finetune_path
      )
      train_state = train_utils.restore_finetune_checkpoint(
          config.vqgan.finetune_path,
          train_state,
          step=None,
          not_restore_g_params_keys=None,
      )

    g_param_size = parameter_overview.count_parameters(train_state.g_params)
    d_param_size = parameter_overview.count_parameters(train_state.d_params)
    writer.write_scalars(
        0, {
            'param_size/generator': g_param_size,
            'param_size/discriminator': d_param_size
        })

  if (start_step == 0  # Which means "no" checkpoint is restored!
      and config.get('init_from') is not None):
    assert not (
        config.get('init_from') and config.get('vqgan.finetune_decoder')
    ), '3d-inflation and finetuning are conflicting setting.'
    init_checkpoint_path = config.init_from.get('checkpoint_path')
    train_state = vqgan_train_lib.init_from_pretrained_checkpoint(
        init_checkpoint_path, train_state, config)

  # Get learning rate scheduler.
  total_steps, steps_per_epoch = scenic_train_utils.get_num_training_steps(
      config, dataset.meta_data)

  train_utils.check_training_step(steps_per_epoch, config)

  # Training step pmap.
  train_state = jax_utils.replicate(train_state)
  train_step_pmapped = jax.pmap(
      functools.partial(
          train_step,
          model_dict=model_dict,
          config=config),
      axis_name='device',
      # We can donate the buffer of train_state.
      donate_argnums=(0),
  )

  sample_step_pmapped = jax.pmap(
      functools.partial(sample_step, model_dict=model_dict, config=config),
      axis_name='device',
      # We can donate the buffer of train_batch.
      donate_argnums=(1),
  )

  checkpoint_steps = config.logging.get('checkpoint_steps', steps_per_epoch)
  log_metric_steps = config.logging.get('log_metric_steps', checkpoint_steps)
  log_sample_size = config.logging.get('log_sample_size', 2)

  logging.info('Starting training loop at step %d of total_steps=%d.',
               start_step, total_steps)

  report_progress = periodic_actions.ReportProgress(
      num_train_steps=total_steps, writer=writer)
  hooks = [report_progress]
  if config.get('xprof', True) and lead_host:
    hooks.append(periodic_actions.Profile(num_profile_steps=5, logdir=workdir))

  train_metrics = None
  with metric_writers.ensure_flushes(writer):
    for step in range(start_step + 1, total_steps + 1):
      with jax.profiler.StepTraceAnnotation('train', step_num=step):
        train_batch = next(dataset.train_iter)
        if config.dataset_configs.get('one_hot_labels', False):
          train_batch.pop('label', None)
        train_state, metrics_update = train_step_pmapped(
            train_state, train_batch)
        metric_update = jax_utils.unreplicate(metrics_update)
        train_metrics = (
            metric_update
            if train_metrics is None else train_metrics.merge(metric_update))

      # Quick indication that training is happening.
      logging.log_first_n(logging.INFO, 'Finished training step %d.', 5, step)
      if step < 200 and config.get('vqgan.finetune_decoder'):
        # checks the finetuning parameter subset.
        encoder_params_sum = train_utils.sum_model_params(
            train_state.g_params['encoder'])
        decoder_params_sum = train_utils.sum_model_params(
            train_state.g_params['decoder'])
        quantizer_params_sum = train_utils.sum_model_params(
            train_state.g_params['quantizer'])
        logging.log_first_n(
            logging.INFO,
            'Params_sum: decoder %f, quantizer %f, encoder %f.', 200,
            decoder_params_sum, quantizer_params_sum, encoder_params_sum)

      for h in hooks:
        try:
          h(step)
        except:  # pylint: disable=bare-except
          logging.exception('Hook failed')
          continue

      ############### LOG TRAIN METRICS ###############
      if (step % log_metric_steps == 1) or (step == total_steps):
        train_metrics = train_metrics.compute()
        train_metrics = {f'train/{k}': v for k, v in train_metrics.items()}
        writer.write_scalars(step, train_metrics)
        writer.write_scalars(step, {'train/lr': learning_rate_fn(step-1)})
        # Reset metric accumulation for next evaluation cycle.
        if step < total_steps:
          train_metrics = None

      ##################### CHECKPOINTING ############################
      if ((step % checkpoint_steps == 1) or (step == total_steps)):
        with report_progress.timed('checkpoint'):
          # Sync model state across replicas.
          train_state = vqgan_train_lib.sync_model_state_across_replicas(
              train_state)
          if config.logging.enable_checkpoint:
            train_utils.save_checkpoint(workdir, train_state,
                                        config.logging.checkpoint_kept)

        ############### LOG TRAIN SAMPLES ###############
        with report_progress.timed('sample'):
          batch_to_sample = jax.tree_util.tree_map(
              lambda x: x[:, :log_sample_size], train_batch)
          batch_samples, _ = sample_step_pmapped(train_state, batch_to_sample)
          batch_samples = jax.device_get(jax_utils.unreplicate(batch_samples))
          log_samples(batch_samples, writer, 'train', step, log_sample_size,
                      config.get('dataset_type', 'video') == 'video')
          del batch_to_sample, batch_samples

  # Tell evaluation job to stop.
  task_manager = eval_utils.TaskManager(workdir)
  task_manager.mark_training_done()
  logging.info('Finishing training at step %d', total_steps)
  # Return the train and eval summary after last step for regression testing.
  return train_state, train_metrics
Main training loop lives in this function.
videogvt/trainers/vqgan_trainer.py