code: def get_estimator_spec(
      hparams,
      mode,
      features,
      labels,
      frame_logits,
      onset_logits,
      offset_logits,
      velocity_values,
      offset_network=True,
  ):
      """Create TPUEstimatorSpec."""
      loss_metrics = {}
      loss = None
      if mode in (
          tf_estimator.ModeKeys.TRAIN,
          tf_estimator.ModeKeys.EVAL,
      ):
          onset_losses = tf.losses.sigmoid_cross_entropy(
              labels.onsets[:, :, : constants.MIDI_PITCHES],
              onset_logits[:, :, : constants.MIDI_PITCHES],
              weights=tf.expand_dims(
                  tf.sequence_mask(
                      features.length, maxlen=tf.shape(labels.onsets)[1]
                  ),
                  axis=2,
              ),
          )
          loss_metrics["onset"] = onset_losses
  
          if offset_network and not hparams.drums_only:
              offset_losses = tf.losses.sigmoid_cross_entropy(
                  labels.offsets[:, :, : constants.MIDI_PITCHES],
                  offset_logits[:, :, : constants.MIDI_PITCHES],
                  weights=tf.expand_dims(
                      tf.sequence_mask(
                          features.length,
                          maxlen=tf.shape(labels.offsets)[1],
                      ),
                      axis=2,
                  ),
              )
              loss_metrics["offset"] = offset_losses
  
          velocity_losses = tf.losses.mean_squared_error(
              labels.velocities,
              velocity_values,
              weights=labels.onsets * hparams.velocity_loss_weight,
          )
          loss_metrics["velocity"] = velocity_losses
  
          if not hparams.drums_only:
              frame_losses = tf.losses.sigmoid_cross_entropy(
                  labels.labels[:, :, : constants.MIDI_PITCHES],
                  frame_logits[:, :, : constants.MIDI_PITCHES],
                  weights=tf.expand_dims(
                      tf.sequence_mask(
                          features.length,
                          maxlen=tf.shape(labels.labels)[1],
                      ),
                      axis=2,
                  ),
              )
              loss_metrics["frame"] = frame_losses
  
          loss = tf.losses.get_total_loss()
  
      if mode in (
          tf_estimator.ModeKeys.EVAL,
          tf_estimator.ModeKeys.PREDICT,
      ):
          frame_probs = tf.sigmoid(frame_logits)
          onset_probs = tf.sigmoid(onset_logits)
          if offset_network:
              offset_probs = tf.sigmoid(offset_logits)
          else:
              offset_probs = tf.zeros_like(onset_probs)
          frame_predictions = (
              frame_probs > hparams.predict_frame_threshold
          )
          onset_predictions = (
              onset_probs > hparams.predict_onset_threshold
          )
          offset_predictions = (
              offset_probs > hparams.predict_offset_threshold
          )
  
          if hparams.drum_prediction_map:
              map_predictions = functools.partial(
                  drum_mappings.map_pianoroll,
                  mapping_name=hparams.drum_prediction_map,
                  reduce_mode="any",
                  min_pitch=constants.MIN_MIDI_PITCH,
              )
              frame_predictions = tf.map_fn(
                  map_predictions, frame_predictions
              )
              onset_predictions = tf.map_fn(
                  map_predictions, onset_predictions
              )
              offset_predictions = tf.map_fn(
                  map_predictions, offset_predictions
              )
              map_values = functools.partial(
                  drum_mappings.map_pianoroll,
                  mapping_name=hparams.drum_prediction_map,
                  reduce_mode="max",
                  min_pitch=constants.MIN_MIDI_PITCH,
              )
              velocity_values = tf.map_fn(map_values, velocity_values)
  
          metrics_values = get_metrics(
              features,
              labels,
              frame_probs,
              onset_probs,
              frame_predictions,
              onset_predictions,
              offset_predictions,
              velocity_values,
              hparams,
          )
  
          for label, loss_collection in loss_metrics.items():
              loss_label = "losses/" + label
              metrics_values[loss_label] = loss_collection
  
      if mode == tf_estimator.ModeKeys.TRAIN:
          train_op = tf_slim.optimize_loss(
              name="training",
              loss=loss,
              global_step=tf.train.get_or_create_global_step(),
              learning_rate=hparams.learning_rate,
              learning_rate_decay_fn=functools.partial(
                  tf.train.exponential_decay,
                  decay_steps=hparams.decay_steps,
                  decay_rate=hparams.decay_rate,
                  staircase=True,
              ),
              clip_gradients=hparams.clip_norm,
              summaries=[],
              optimizer=lambda lr: tf.tpu.CrossShardOptimizer(
                  tf.train.AdamOptimizer(lr)
              ),
          )
  
          return tf.tpu.estimator.TPUEstimatorSpec(
              mode=mode, loss=loss, train_op=train_op
          )
      elif mode == tf_estimator.ModeKeys.EVAL:
          metric_ops = {
              k: tf.metrics.mean(v) for k, v in metrics_values.items()
          }
          return tf_estimator.EstimatorSpec(
              mode=mode, loss=loss, eval_metric_ops=metric_ops
          )
      elif mode == tf_estimator.ModeKeys.PREDICT:
          predictions = {
              "frame_probs": frame_probs,
              "onset_probs": onset_probs,
              "frame_predictions": frame_predictions,
              "onset_predictions": onset_predictions,
              "offset_predictions": offset_predictions,
              "velocity_values": velocity_values,
              "sequence_predictions": _predict_sequences(
                  frame_probs=frame_probs,
                  onset_probs=onset_probs,
                  frame_predictions=frame_predictions,
                  onset_predictions=onset_predictions,
                  offset_predictions=offset_predictions,
                  velocity_values=velocity_values,
                  hparams=hparams,
              ),
              # Include some features and labels in output because Estimator 'predict'
              # API does not give access to them.
              "sequence_ids": features.sequence_id,
              "sequence_labels": labels.note_sequence,
              "frame_labels": labels.labels,
              "onset_labels": labels.onsets,
          }
          for k, v in metrics_values.items():
              predictions[k] = tf.stack(v)
  
          return tf_estimator.EstimatorSpec(
              mode=mode, predictions=predictions
          )
      else:
          raise ValueError("Unsupported mode: %s" % mode)
  
file_path: magenta/magenta/models/onsets_frames_transcription/estimator_spec_util.py
imports: import functools
  from magenta.models.onsets_frames_transcription import constants
  from magenta.models.onsets_frames_transcription import drum_mappings
  from magenta.models.onsets_frames_transcription import infer_util
  from magenta.models.onsets_frames_transcription import metrics
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  import tf_slim
label_desc: Use image node_id 4 for calling a global function with example usage: get_estimator_spec(hparams, mode, features, labels, frame_logits, onset_logits, offset_logits, velocity_values, offset_network) and returns: tf, tf_estimator, tf_estimator