code: def train(
      data_provider,
      trainer,
      batch_size=32,
      num_steps=1000000,
      steps_per_summary=300,
      steps_per_save=300,
      save_dir="/tmp/ddsp",
      restore_dir="/tmp/ddsp",
      early_stop_loss_value=None,
      report_loss_to_hypertune=False,
  ):
      """Main training loop.
  
      Args:
       data_provider: DataProvider object for training data.
       trainer: Trainer object built with Model to train.
       batch_size: Total batch size.
       num_steps: Number of training steps.
       steps_per_summary: Number of training steps per summary save.
       steps_per_save: Number of training steps per checkpoint save.
       save_dir: Directory where checkpoints and summaries will be saved.
         If empty string, no checkpoints or summaries will be saved.
       restore_dir: Directory where latest checkpoints for resuming the training
         are stored. If there are no checkpoints in this directory, training will
         begin anew.
       early_stop_loss_value: Early stopping. When the total_loss reaches below this
         value training stops. If None training will run for num_steps steps.
       report_loss_to_hypertune: Report loss values to hypertune package for
         hyperparameter tuning, such as on Google Cloud AI-Platform.
      """
      # Get a distributed dataset iterator.
      dataset = data_provider.get_batch(
          batch_size, shuffle=True, repeats=-1
      )
      dataset = trainer.distribute_dataset(dataset)
      dataset_iter = iter(dataset)
  
      # Build model, easiest to just run forward pass.
      trainer.build(next(dataset_iter))
  
      # Load latest checkpoint if one exists in load directory.
      try:
          trainer.restore(restore_dir)
      except FileNotFoundError:
          logging.info(
              "No existing checkpoint found in %s, skipping "
              "checkpoint loading.",
              restore_dir,
          )
  
      if save_dir:
          # Set up the summary writer and metrics.
          summary_dir = os.path.join(save_dir, "summaries", "train")
          summary_writer = tf.summary.create_file_writer(summary_dir)
  
          # Save the gin config.
          write_gin_config(
              summary_writer, save_dir, trainer.step.numpy()
          )
      else:
          # Need to create a dummy writer, even if no save_dir is provided.
          summary_writer = tf.summary.create_noop_writer()
  
      # Train.
      with summary_writer.as_default():
          tick = time.time()
  
          first_step = True
  
          while trainer.step < num_steps:
              step = trainer.step
  
              # Take a step.
              losses = trainer.train_step(dataset_iter)
  
              # Create training loss metrics when starting/restarting training.
              if first_step:
                  loss_names = list(losses.keys())
                  logging.info("Creating metrics for %s", loss_names)
                  avg_losses = {
                      name: tf.keras.metrics.Mean(
                          name=name, dtype=tf.float32
                      )
                      for name in loss_names
                  }
                  first_step = False
  
              # Update metrics.
              for k, v in losses.items():
                  avg_losses[k].update_state(v)
  
              # Log the step.
              log_str = "step: {}\t".format(int(step.numpy()))
              for k, v in losses.items():
                  log_str += "{}: {:.2f}\t".format(k, v)
              logging.info(log_str)
  
              # Write Summaries.
              if step % steps_per_summary == 0 and save_dir:
                  # Speed.
                  steps_per_sec = steps_per_summary / (
                      time.time() - tick
                  )
                  tf.summary.scalar(
                      "steps_per_sec", steps_per_sec, step=step
                  )
                  tick = time.time()
  
                  # Metrics.
                  for k, metric in avg_losses.items():
                      tf.summary.scalar(
                          "losses/{}".format(k),
                          metric.result(),
                          step=step,
                      )
                      metric.reset_states()
  
              # Report metrics for hyperparameter tuning if enabled.
              if report_loss_to_hypertune:
                  cloud.report_metric_to_hypertune(
                      losses["total_loss"], step.numpy()
                  )
  
              # Stop the training when the loss reaches given value
              if (
                  early_stop_loss_value is not None
                  and losses["total_loss"] <= early_stop_loss_value
              ):
                  logging.info(
                      "Total loss reached early stopping value of %s",
                      early_stop_loss_value,
                  )
                  break
  
              # Save Model.
              if step % steps_per_save == 0 and save_dir:
                  trainer.save(save_dir)
                  summary_writer.flush()
  
      # Write a final checkpoint.
      if save_dir:
          trainer.save(save_dir)
          summary_writer.flush()
  
      logging.info("Training Finished!")
  
file_path: ddsp/ddsp/training/train_util.py
imports: import json
  import os
  import time
  from absl import logging
  from ddsp.training import cloud
  import gin
  import tensorflow.compat.v2
label_desc: Use image node_id 7 for calling a global function with example usage: train(data_provider, trainer, batch_size, num_steps, steps_per_summary, steps_per_save, save_dir, restore_dir, early_stop_loss_value, report_loss_to_hypertune) without return types