code: def train(model_func, params):
      image_width = params["image_width"]
      image_height = params["image_height"]
      image_format = params["image_format"]
      distort_color = params["distort_color"]
      momentum = params["momentum"]
      loss_scale = params["loss_scale"]
      data_dir = params["data_dir"]
      data_idx_dir = params["data_idx_dir"]
      batch_size = params["batch_size"]
      num_iter = params["num_iter"]
      iter_unit = params["iter_unit"]
      log_dir = params["log_dir"]
      export_dir = params["export_dir"]
      tensorboard_dir = params["tensorboard_dir"]
      display_every = params["display_every"]
      precision = params["precision"]
      dali_mode = params["dali_mode"]
      use_xla = params["use_xla"]
  
      if data_dir is not None:
          file_format = os.path.join(data_dir, "%s-*")
          train_files = sorted(tf.io.gfile.glob(file_format % "train"))
          valid_files = sorted(
              tf.io.gfile.glob(file_format % "validation")
          )
          num_train_samples = common.get_num_records(train_files)
          num_valid_samples = common.get_num_records(valid_files)
      else:
          num_train_samples = 1281982
          num_valid_samples = 5000
  
      train_idx_files = None
      valid_idx_files = None
      if data_idx_dir is not None:
          file_format = os.path.join(data_idx_dir, "%s-*")
          train_idx_files = sorted(
              tf.io.gfile.glob(file_format % "train")
          )
          valid_idx_files = sorted(
              tf.io.gfile.glob(file_format % "validation")
          )
  
      if iter_unit.lower() == "epoch":
          num_epochs = num_iter
          nstep_per_epoch = num_train_samples // (
              batch_size * hvd.size()
          )
          nstep_per_valid = num_valid_samples // (
              batch_size * hvd.size()
          )
      else:
          assert iter_unit.lower() == "batch"
          num_epochs = 1
          nstep_per_epoch = min(
              num_iter, num_train_samples // (batch_size * hvd.size())
          )
          nstep_per_valid = min(
              10, num_valid_samples // (batch_size * hvd.size())
          )
  
      initial_epoch = 0
      if log_dir:
          # We save check points only when using the real data.
          assert (
              data_dir
          ), "--data_dir cannot be empty when using --log_dir"
          assert os.path.exists(log_dir)
          ckpt_format = (
              log_dir + "/model-{epoch:02d}-{val_top1:.2f}.hdf5"
          )
          # Looks for the most recent checkpoint and sets the initial epoch from it.
          for filename in os.listdir(log_dir):
              if filename.startswith("model-"):
                  initial_epoch = max(
                      int(re.findall(r"\d+", filename)[0]),
                      initial_epoch,
                  )
  
      if tensorboard_dir:
          assert os.path.exists(tensorboard_dir)
  
      if export_dir:
          assert os.path.exists(export_dir)
          save_format = export_dir + "/saved_model_rn50.h5"
  
      if use_xla:
          tf.config.optimizer.set_jit(True)
  
      # Horovod: pin GPU to be used to process local rank (one GPU per process)
      gpus = tf.config.experimental.list_physical_devices("GPU")
      for gpu in gpus:
          tf.config.experimental.set_memory_growth(gpu, True)
      if gpus:
          tf.config.experimental.set_visible_devices(
              gpus[hvd.local_rank()], "GPU"
          )
  
      if precision == "fp16":
          if StrictVersion(tf.__version__) >= StrictVersion("2.4.0"):
              policy = keras.mixed_precision.Policy("mixed_float16")
              keras.mixed_precision.set_global_policy(policy)
          else:
              policy = keras.mixed_precision.experimental.Policy(
                  "mixed_float16", loss_scale
              )
              keras.mixed_precision.experimental.set_policy(policy)
  
      lr_schedule = common.create_piecewise_constant_decay_with_warmup(
          batch_size=batch_size * hvd.size(),
          epoch_size=num_train_samples,
          warmup_epochs=common.LR_SCHEDULE[0][1],
          boundaries=list(p[1] for p in common.LR_SCHEDULE[1:]),
          multipliers=list(p[0] for p in common.LR_SCHEDULE),
          compute_lr_on_cpu=True,
      )
      opt = keras.optimizers.SGD(
          learning_rate=lr_schedule, momentum=momentum
      )
      # Horovod: add Horovod DistributedOptimizer. We use a modified version to
      # support the custom learning rate schedule.
      opt = hvd.DistributedOptimizer(opt)
      if (
          StrictVersion(tf.__version__) >= StrictVersion("2.4.0")
          and precision == "fp16"
      ):
          opt = keras.mixed_precision.LossScaleOptimizer(
              opt, dynamic=False, initial_scale=loss_scale
          )
  
      backend.set_image_data_format(image_format)
      dtype = "float16" if precision == "fp16" else "float32"
      backend.set_floatx(dtype)
      model = model_func(num_classes=image_processing.NUM_CLASSES)
      loss_func = ("sparse_categorical_crossentropy",)
  
      top5 = tf.keras.metrics.SparseTopKCategoricalAccuracy(
          k=5, name="top5"
      )
      top1 = tf.keras.metrics.SparseTopKCategoricalAccuracy(
          k=1, name="top1"
      )
  
      # Horovod: Specify `experimental_run_tf_function=False` to ensure TensorFlow
      # uses hvd.DistributedOptimizer() to compute gradients. However, this option
      # will disable the overlapping of the data loading and compute and hurt the
      # performace if the model is not under the scope of distribution strategy
      # scope.
      model.compile(
          optimizer=opt,
          loss=loss_func,
          metrics=[top1, top5],
          experimental_run_tf_function=False,
      )
  
      training_hooks = []
      training_hooks.append(
          hvd.callbacks.BroadcastGlobalVariablesCallback(0)
      )
      training_hooks.append(
          _ProfileKerasFitCallback(batch_size, display_every)
      )
  
      if log_dir and hvd.rank() == 0:
          ckpt_callback = keras.callbacks.ModelCheckpoint(
              ckpt_format,
              monitor="val_top1",
              verbose=1,
              save_best_only=False,
              save_weights_only=False,
              save_frequency=1,
          )
          training_hooks.append(ckpt_callback)
  
      if tensorboard_dir and hvd.rank() == 0:
          tensorboard_callback = tf.keras.callbacks.TensorBoard(
              log_dir=tensorboard_dir
          )
          training_hooks.append(tensorboard_callback)
  
      if data_dir is not None:
          num_preproc_threads = (
              params["dali_threads"] if dali_mode else 10
          )
          train_input = image_processing.image_set(
              train_files,
              batch_size,
              image_height,
              image_width,
              training=True,
              distort_color=distort_color,
              deterministic=False,
              num_threads=num_preproc_threads,
              use_dali=dali_mode,
              idx_filenames=train_idx_files,
          )
  
          valid_input = image_processing.image_set(
              valid_files,
              batch_size,
              image_height,
              image_width,
              training=False,
              distort_color=False,
              deterministic=False,
              num_threads=num_preproc_threads,
              use_dali=dali_mode,
              idx_filenames=valid_idx_files,
          )
          if dali_mode:
              train_input = train_input.get_device_dataset()
              valid_input = valid_input.get_device_dataset()
          valid_params = {
              "validation_data": valid_input,
              "validation_steps": nstep_per_valid,
              "validation_freq": 1,
          }
      else:
          train_input = image_processing.fake_image_set(
              batch_size, image_height, image_width
          )
          valid_params = {}
  
      try:
          verbose = 2 if hvd.rank() == 0 else 0
          model.fit(
              train_input,
              epochs=num_epochs,
              callbacks=training_hooks,
              steps_per_epoch=nstep_per_epoch,
              verbose=verbose,
              initial_epoch=initial_epoch,
              **valid_params,
          )
      except KeyboardInterrupt:
          print("Keyboard interrupt")
  
      if export_dir and hvd.rank() == 0:
          model.save(save_format)
          print(f"The model is saved to {save_format}")
  
file_path: DALI/docs/examples/use_cases/tensorflow/resnet-n/nvutils/runner.py
imports: from nvutils import image_processing
  from nvutils import common
  from distutils.version import StrictVersion
  import tensorflow
  import keras
  import os
  import time
  import re
  import horovod.tensorflow.keras
  from keras import backend
label_desc: Use image node_id 1 for calling a global function with example usage: train(model_func, params) without return types