code: def setup_keras_trainer(
      trainer,
      get_model,
      input_signature,
      target_signature,
      input,
      optimizer,
      loss,
      metrics,
  ):
      """
      Args:
          trainer (SingleCostTrainer):
          get_model (input1, input2, ... -> tf.keras.Model):
              A function which takes tensors, builds and returns a Keras model.
              It will be part of the tower function.
          input (InputSource):
          optimizer (tf.train.Optimizer):
          loss, metrics: list of strings
      """
      assert isinstance(optimizer, tf.train.Optimizer), optimizer
      assert isinstance(loss, list), loss
      assert len(loss) >= 1, "No loss was given!"
      assert isinstance(metrics, list), metrics
      model_caller = KerasModelCaller(get_model)
  
      nr_inputs = len(input_signature)
  
      def get_cost(*inputs):
          ctx = get_current_tower_context()
          input_tensors = list(inputs[:nr_inputs])
          target_tensors = list(inputs[nr_inputs:])
          # TODO mapping between target tensors & output tensors
  
          outputs = model_caller(*input_tensors)
  
          if isinstance(outputs, tf.Tensor):
              outputs = [outputs]
          assert len(outputs) == len(
              target_tensors
          ), "len({}) != len({})".format(
              str(outputs), str(target_tensors)
          )
          assert len(outputs) == len(loss), "len({}) != len({})".format(
              str(outputs), str(loss)
          )
  
          loss_tensors = []
          for idx, loss_name in enumerate(loss):
              with cached_name_scope("keras_loss", top_level=False):
                  loss_fn = keras.losses.get(loss_name)
                  curr_loss = loss_fn(target_tensors[idx], outputs[idx])
              curr_loss = tf.reduce_mean(curr_loss, name=loss_name)
              _check_name(curr_loss, loss_name)
              loss_tensors.append(curr_loss)
  
          loss_reg = regularize_cost_from_collection()
          if loss_reg is not None:
              total_loss = tf.add_n(
                  loss_tensors + [loss_reg], name=TOTAL_LOSS_NAME
              )
              add_moving_summary(loss_reg, total_loss, *loss_tensors)
          else:
              total_loss = tf.add_n(loss_tensors, name=TOTAL_LOSS_NAME)
              add_moving_summary(total_loss, *loss_tensors)
  
          if metrics and (
              ctx.is_main_training_tower or not ctx.is_training
          ):
              # for list: one metric for each output
              metric_tensors = []
              for oid, metric_name in enumerate(metrics):
                  output_tensor = outputs[oid]
                  target_tensor = target_tensors[
                      oid
                  ]  # TODO may not have the same mapping?
                  with cached_name_scope(
                      "keras_metric", top_level=False
                  ):
                      metric_fn = keras.metrics.get(metric_name)
                      metric_tensor = metric_fn(
                          target_tensor, output_tensor
                      )
                  metric_tensor = tf.reduce_mean(
                      metric_tensor, name=metric_name
                  )
                  _check_name(metric_tensor, metric_name)
                  # check name conflict here
                  metric_tensors.append(metric_tensor)
              add_moving_summary(*metric_tensors)
  
          return total_loss
  
      trainer.setup_graph(
          input_signature + target_signature,
          input,
          get_cost,
          lambda: optimizer,
      )
      if (
          isinstance(keras.backend.learning_phase(), tf.Tensor)
          and len(keras.backend.learning_phase().consumers()) > 0
      ):
          # check if learning_phase is used in this model
          trainer.register_callback(KerasPhaseCallback(True))
  
imports: from contextlib import contextmanager
  import six
  import tensorflow
  from tensorflow import keras
  from ..callbacks import Callback, CallbackToHook, InferenceRunner, InferenceRunnerBase, ScalarStats
  from ..models.regularize import regularize_cost_from_collection
  from ..tfutils.collection import backup_collection, restore_collection
  from ..tfutils.common import get_op_tensor_name
  from ..tfutils.scope_utils import cached_name_scope
  from ..tfutils.summary import add_moving_summary
  from ..tfutils.tower import get_current_tower_context
  from ..train import SimpleTrainer, SyncMultiGPUTrainerParameterServer, Trainer
  from ..train.interface import apply_default_prefetch
  from ..train.trainers import DistributedTrainerBase
  from ..utils import logger
  from ..utils.gpu import get_nr_gpu
label_desc: Use image node_id 2 for calling a global function with example usage: setup_keras_trainer(trainer, get_model, input_signature, target_signature, input, optimizer, loss, metrics) and returns: total_loss