code: def create_estimator(
      model_fn,
      model_dir,
      hparams,
      use_tpu=False,
      master="",
      tpu_cluster=None,
      save_checkpoint_steps=300,
      save_summary_steps=300,
      keep_checkpoint_max=None,
      warm_start_from=None,
  ):
      """Creates an estimator."""
  
      def wrapped_model_fn(features, labels, mode, params, config):
          """Wrap model_fn to restore labels value if present in features."""
          # Workaround for Estimator API that forces 'labels' to be None when in
          # predict mode.
          # https://github.com/tensorflow/tensorflow/issues/17824
          # See also infer_util.labels_to_features_wrapper
          if labels is None and hasattr(features, "labels"):
              labels = features.labels
          return model_fn(features, labels, mode, params, config)
  
      if tpu_cluster:
          tpu_cluster_resolver = (
              tf.distribute.cluster_resolver.TPUClusterResolver(
                  tpu_cluster
              )
          )
          master = None
      else:
          tpu_cluster_resolver = None
  
      config = tf_estimator.tpu.RunConfig(
          tpu_config=tf_estimator.tpu.TPUConfig(
              iterations_per_loop=save_checkpoint_steps
          ),
          master=master,
          cluster=tpu_cluster_resolver,
          save_summary_steps=save_summary_steps,
          save_checkpoints_steps=save_checkpoint_steps,
          keep_checkpoint_max=keep_checkpoint_max,
          keep_checkpoint_every_n_hours=1,
      )
  
      params = copy.deepcopy(hparams)
      params.del_hparam("batch_size")
      return tf_estimator.tpu.TPUEstimator(
          use_tpu=use_tpu,
          model_fn=wrapped_model_fn,
          model_dir=model_dir,
          params=params,
          train_batch_size=hparams.batch_size,
          eval_batch_size=hparams.eval_batch_size,
          predict_batch_size=hparams.predict_batch_size,
          config=config,
          warm_start_from=warm_start_from,
          eval_on_tpu=False,
      )
  
file_path: magenta/magenta/models/onsets_frames_transcription/train_util.py
imports: import collections
  import copy
  import functools
  import random
  import sys
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  import tf_slim
label_desc: Use image node_id 2 for calling a global function with example usage: create_estimator(model_fn, model_dir, hparams, use_tpu, master, tpu_cluster, save_checkpoint_steps, save_summary_steps, keep_checkpoint_max, warm_start_from) and returns: tf_estimator, model_fn