code: def train():
      """Train a model."""
      if FLAGS.train_config is None:
          raise ValueError("train_config must be specified")
  
      experiment_dir = _EXPERIMENT_DIR.value
      experiment_dir = os.path.join(
          experiment_dir, construct_hyperparam_str()
      )
      experiment_model_dir = f"{experiment_dir}/checkpoints"
      experiment_tensorboard_dir = f"{experiment_dir}/tensorboard_log"
      experiment_backup_dir = f"{experiment_dir}/backup"
      logging.info(
          "Use TPU at %s",
          _TPU.value if _TPU.value is not None else "local",
      )
      logging.info("experiment_dir: %s", experiment_dir)
      tpu = _TPU.value
      if tpu is not None:
          resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
              tpu=tpu
          )
          tf.config.experimental_connect_to_cluster(resolver)
          tf.tpu.experimental.initialize_tpu_system(resolver)
          strategy = tf.distribute.TPUStrategy(resolver)
      else:
          strategy = tf.distribute.MirroredStrategy()
  
      train_dataset_config = data_providers.read_dataset_config(
          FLAGS.train_config.train_dataset_pbtxt
      )
  
      eval_dataset_config = data_providers.read_dataset_config(
          FLAGS.train_config.tune_dataset_pbtxt
      )
      input_shape = dv_utils.get_shape_from_examples_path(
          train_dataset_config.tfrecord_path
      )
      # Set the total number epochs. This is the number of times we will iterate
      # over the data.
      epochs = FLAGS.train_config.num_epochs
      # Calculate the total number of mini_epochs, we use it to report metrics
      # sooner. This divides a long waited epoch into smaller epochs. This is the
      # value fed into model.fit so it is assumed to be one epoch when it reaches
      # a mini epoch.
      mini_epochs = (
          epochs * FLAGS.train_config.num_mini_epochs_per_epoch
      )
      # Steps per epoch is the total number of steps/batches we have in each epoch.
      steps_per_epoch = (
          train_dataset_config.num_examples
          // FLAGS.train_config.batch_size_train
      )
      # Steps per mini-epoch is after how many batches we report a metric. This is
      # what used in model.fit so it reports metrics after steps_per_mini_epoch
      # numbers.
      steps_per_mini_epoch = (
          steps_per_epoch
          // FLAGS.train_config.num_mini_epochs_per_epoch
      )
      if (
          FLAGS.train_config.denovo_enabled
          and FLAGS.train_config.denovo_weight == 1.0
      ):
          logging.warning(
              "Denovo is enabled but denovo weight is set to 1.0"
          )
  
      logging.info(
          (
              "1) Total examples for training: %s.\n"
              "2) Train batch size or examples per step: %s.\n"
              "3) Total epochs: %s.\n"
              "4) Total steps per epoch: %s.\n"
              "5) Number of mini-epochs per epoch: %s.\n"
              "6) Total steps per mini-epoch: %s.\n"
              "7) Total mini-epochs to training completion (x-axis in xm): %s.\n"
              "8) Denovo weight: %s.\n"
              "9) Denovo enabled: %s. \b"
          ),
          train_dataset_config.num_examples,
          FLAGS.train_config.batch_size_train,
          epochs,
          steps_per_epoch,
          FLAGS.train_config.num_mini_epochs_per_epoch,
          steps_per_mini_epoch,
          epochs * FLAGS.train_config.num_mini_epochs_per_epoch,
          FLAGS.train_config.denovo_weight,
          FLAGS.train_config.denovo_enabled,
      )
  
      with strategy.scope():
          init_weights_file = None
          if FLAGS.train_config.init_weights_file:
              init_weights_file = FLAGS.train_config.init_weights_file
          model = modeling.inceptionv3(input_shape, init_weights_file)
  
          decay_steps = int(
              steps_per_epoch
              * FLAGS.train_config.learning_rate_num_epochs_per_decay
          )
  
          learning_rate = tf.keras.optimizers.schedules.ExponentialDecay(
              initial_learning_rate=FLAGS.train_config.learning_rate,
              decay_steps=decay_steps,
              decay_rate=FLAGS.train_config.learning_rate_decay_rate,
              staircase=True,
          )
  
          if FLAGS.train_config.warmup_mini_epochs > 0:
              warmup_steps = int(
                  steps_per_mini_epoch
                  * FLAGS.train_config.warmup_mini_epochs
              )
              warmup_learning_rate = (
                  FLAGS.train_config.learning_rate / 10
              )
              logging.info(
                  (
                      "Use LinearWarmup. warmup_mini_epochs=%s, warmup_steps=%s, "
                      "warmup_learning_rate=%s"
                  ),
                  FLAGS.train_config.warmup_mini_epochs,
                  warmup_steps,
                  warmup_learning_rate,
              )
              learning_rate = optimization.LinearWarmup(
                  # This is initial learning rate.
                  warmup_learning_rate=warmup_learning_rate,
                  after_warmup_lr_sched=learning_rate,
                  warmup_steps=warmup_steps,
              )
  
          # Our "steps_per_epoch" in model.fit is:
          #   steps_per_epoch // FLAGS.train_config.num_mini_epochs_per_epoch.
          # I divided that by 10 just to be safe. In case that gets too small,
          # I set a max with 128 here.
          # Note that this might only make sense for TPUs. In the future we'll want to
          # check whether this works for GPU or not.
          # We should also keep an eye on this feature which might allow autotune in
          # the future: https://github.com/keras-team/keras/issues/16573.
          steps_per_execution = max(
              128,
              steps_per_epoch
              // FLAGS.train_config.num_mini_epochs_per_epoch
              // 10,
          )
          logging.info("steps_per_execution=%s", steps_per_execution)
          model.compile(
              # TODO: Try different optimizers
              # and sweep over these hyperparams
              optimizer=tfa.optimizers.MovingAverage(
                  optimizer=tf.keras.optimizers.RMSprop(
                      learning_rate=learning_rate,
                      rho=FLAGS.train_config.rho,
                      momentum=FLAGS.train_config.momentum,
                      epsilon=FLAGS.train_config.epsilon,
                  ),
                  average_decay=FLAGS.train_config.average_decay,
              ),
              # This is from:
              # https://www.tensorflow.org/guide/tpu#train_the_model_using_keras_high-level_apis.
              # Anything between 2 and `steps_per_epoch` could help here.
              steps_per_execution=steps_per_execution,
              # TODO: Sweep over these hyperparams
              # TODO: Does computing fewer TensorBoard metrics speed up
              # training?
              loss=tf.keras.losses.CategoricalCrossentropy(
                  label_smoothing=FLAGS.train_config.label_smoothing
              ),
              metrics=[
                  tf.keras.metrics.CategoricalAccuracy(),
                  tf.keras.metrics.CategoricalCrossentropy(),
                  tf.keras.metrics.TruePositives(),
                  tf.keras.metrics.TrueNegatives(),
                  tf.keras.metrics.FalsePositives(),
                  tf.keras.metrics.FalseNegatives(),
                  tf.keras.metrics.Precision(),
                  tf.keras.metrics.Precision(
                      name="precision_homref", class_id=0
                  ),
                  tf.keras.metrics.Precision(
                      name="precision_het", class_id=1
                  ),
                  tf.keras.metrics.Precision(
                      name="precision_homalt", class_id=2
                  ),
                  tf.keras.metrics.Recall(),
                  tf.keras.metrics.Recall(
                      name="recall_homref", class_id=0
                  ),
                  tf.keras.metrics.Recall(
                      name="recall_het", class_id=1
                  ),
                  tf.keras.metrics.Recall(
                      name="recall_homalt", class_id=2
                  ),
                  tfa.metrics.F1Score(
                      num_classes=3,
                      average="weighted",
                      name="f1_weighted",
                  ),
                  tfa.metrics.F1Score(
                      num_classes=3, average="micro", name="f1_micro"
                  ),
                  tfa.metrics.F1Score(
                      num_classes=3, average="macro", name="f1_macro"
                  ),
                  F1ScorePerClass(
                      num_classes=3, target_class=0, name="f1_homref"
                  ),
                  F1ScorePerClass(
                      num_classes=3, target_class=1, name="f1_het"
                  ),
                  F1ScorePerClass(
                      num_classes=3, target_class=2, name="f1_homalt"
                  ),
              ],
              weighted_metrics=[
                  tf.keras.metrics.CategoricalAccuracy(
                      name="weighted_categorical_accuracy"
                  ),
                  tf.keras.metrics.CategoricalCrossentropy(
                      name="weighted_categorical_crossentropy"
                  ),
                  tf.keras.metrics.TruePositives(
                      name="weighted_true_positives"
                  ),
                  tf.keras.metrics.TrueNegatives(
                      name="weighted_true_negatives"
                  ),
                  tf.keras.metrics.FalsePositives(
                      name="weighted_false_positives"
                  ),
                  tf.keras.metrics.FalseNegatives(
                      name="weighted_false_negatives"
                  ),
                  tf.keras.metrics.Precision(name="weighted_precision"),
              ],
          )
  
      tensorboard_callback = tf.keras.callbacks.TensorBoard(
          log_dir=experiment_tensorboard_dir,
          write_images=False,
          profile_batch=0,
      )
  
      checkpoint_file_avg = os.path.join(
          experiment_model_dir,
          "weights-{epoch:02d}-{%s:f}.ckpt"
          % FLAGS.train_config.best_metrics,
      )
      logging.info(
          "Metrics used to choose best ckpt is: %s",
          FLAGS.train_config.best_metrics,
      )
      logging.info("Best ckpt file path: %s", checkpoint_file_avg)
      logging.info("save_best_only: %s", _SAVE_BEST_ONLY.value)
  
      checkpoint_path = os.path.join(experiment_model_dir, "checkpoint")
      best_metrics_from_checkpoint = None
      weights_path_from_checkpoint = None
      if tf.io.gfile.exists(checkpoint_path):
          with tf.io.gfile.GFile(checkpoint_path, "r") as ckpt_file:
              first_line = ckpt_file.readline()
              match = re.search(
                  r'model_checkpoint_path: "(.*)"', first_line
              )
              if match is None:
                  raise ValueError("Unexpected checkpoint format.")
              weights_path_from_checkpoint = match.group(1)
              match = re.search(
                  r"weights-\d+-(0.\d+).ckpt",
                  weights_path_from_checkpoint,
              )
              if match is None:
                  raise ValueError("Unexpected checkpoint format.")
              match = re.search(
                  r'model_checkpoint_path: "weights-\d+-(0.\d+).ckpt"',
                  first_line,
              )
              if match is None:
                  raise ValueError("Unexpected checkpoint format.")
              best_metrics_from_checkpoint = float(match.group(1))
              logging.info(
                  "Resuming: Loading best metrics from %s: %s",
                  checkpoint_path,
                  best_metrics_from_checkpoint,
              )
      average_model_checkpoint_callback = AverageModelCheckpointPatched(
          filepath=str(checkpoint_file_avg),
          update_weights=True,
          save_best_only=_SAVE_BEST_ONLY.value,
          save_weights_only=True,
          monitor=FLAGS.train_config.best_metrics,
          mode="max",
          save_freq="epoch",
          best_metrics_from_checkpoint=best_metrics_from_checkpoint,
      )
      if weights_path_from_checkpoint is not None:
          weights_path_from_checkpoint = os.path.join(
              experiment_model_dir, weights_path_from_checkpoint
          )
          logging.info(
              "Resuming: load_weights from from %s",
              weights_path_from_checkpoint,
          )
          model.load_weights(weights_path_from_checkpoint)
  
      backup_callback = tf.keras.callbacks.BackupAndRestore(
          experiment_backup_dir
      )
  
      # XManager callback.
      xmanager_callback = XManagerCallback()
  
      training_callbacks = [
          tensorboard_callback,
          average_model_checkpoint_callback,
          backup_callback,
          xmanager_callback,
          ClearSessionCallback(),
          tf.keras.callbacks.EarlyStopping(
              monitor=FLAGS.train_config.best_metrics,
              # `patient` below might need to be tuned.
              patience=FLAGS.train_config.early_stopping.patience
              * FLAGS.train_config.num_mini_epochs_per_epoch,
              mode="max",
              verbose=1,
          ),
      ]
      # Calculate validation attributes.
      num_validation_examples = eval_dataset_config.num_examples
      if FLAGS.train_config.num_validation_examples:
          logging.info(
              (
                  "train_config.num_validation_examples is set. Using "
                  "min(%s, %s) as num_validation_examples."
              ),
              FLAGS.train_config.num_validation_examples,
              eval_dataset_config.num_examples,
          )
          num_validation_examples = min(
              FLAGS.train_config.num_validation_examples,
              eval_dataset_config.num_examples,
          )
      logging.info(
          "num_validation_examples = %s", num_validation_examples
      )
      validation_steps = (
          num_validation_examples // FLAGS.train_config.batch_size_eval
      )
      logging.info("validation_steps: %s", validation_steps)
      model.fit(
          input_fn(
              train_dataset_config.tfrecord_path,
              input_shape,
              mode="train",
          ),
          epochs=mini_epochs,
          verbose=2,
          steps_per_epoch=steps_per_mini_epoch,
          callbacks=training_callbacks,
          validation_data=input_fn(
              eval_dataset_config.tfrecord_path,
              input_shape,
              max_examples=num_validation_examples,
              mode="validation",
          ),
          # Even with "mini epoch", we still want to evaluate the same amount of
          # validation examples per point. So, I'm not dividing this by
          # FLAGS.train_config.num_mini_epochs_per_epoch
          validation_steps=validation_steps,
      )
  
imports: import gc
  import os
  import re
  from typing import Any, Dict, Optional, Tuple
  import warnings
  from absl import app
  from absl import flags
  from absl import logging
  from ml_collections.config_flags import config_flags
  import tensorflow
  import tensorflow_addons
  from deepvariant import data_providers
  from deepvariant import dv_constants
  from deepvariant import dv_utils
  from deepvariant import keras_modeling
  from deepvariant.average_model_checkpoint_patched import AverageModelCheckpointPatched
  from third_party.nucleus.io import sharded_file_utils
  from official.modeling import optimization
label_desc: Use image node_id 5 for calling a global function with example usage: train() without return types