code: def train_and_eval_model(
      model_clf: extended_model.FeatureExtractorWithClassifier,
      train_dataset: tf.distribute.DistributedDataset,
      train_steps_per_epoch: int,
      test_dataset: tf.distribute.DistributedDataset,
      eval_steps_per_epoch: int,
      optimizer: tf.keras.optimizers.Optimizer,
      strategy: tf.distribute.Strategy,
  ) -> None:
      """Train and evaluate the embeddings and the classifier on top them.
  
      Args:
        model_clf: Model to train.
        train_dataset: Dataset over which to perform training. This dataset is
          comprised of tf.train.Example records, containing features 'image' and
          'label', corresponding to the obfuscated image and its respective label.
        train_steps_per_epoch: Number of training steps per epoch.
        test_dataset: Dataset over which to evaluate the model.
        eval_steps_per_epoch: Number of evaluation steps per epoch.
        optimizer: Optimizer to use.
        strategy: Strategy to be used for distributed training.
      """
      with strategy.scope():
          loss_names = [
              "clf_train_loss",
              "clf_train_classification_loss",
              "clf_train_reg_loss",
              "clf_train_embed_loss",
          ]
          for obf in _OBFUSCATIONS_EVAL.value:
              loss_names.append("clf_test_loss_{}".format(obf))
          loss_metrics = {
              loss_name: tf.keras.metrics.Mean(
                  loss_name, dtype=tf.float32
              )
              for loss_name in loss_names
          }
  
          acc_names = ["clf_train_acc"]
          for obf in _OBFUSCATIONS_EVAL.value:
              acc_names.append("clf_test_acc_{}".format(obf))
          acc_metrics = {
              acc_name: tf.keras.metrics.SparseCategoricalAccuracy(
                  acc_name, dtype=tf.float32
              )
              for acc_name in acc_names
          }
  
          clf_checkpoint = tf.train.Checkpoint(
              model=model_clf, optimizer=optimizer
          )
  
      clf_train_summary_writer = tf.summary.create_file_writer(
          os.path.join(
              _MODEL_DIR.value,
              "summaries"
              f"lr_{_BASE_LR.value}_decay_{_WEIGHT_DECAY.value}",
              "clf_train",
          )
      )
  
      clf_test_summary_writer = tf.summary.create_file_writer(
          os.path.join(
              _MODEL_DIR.value,
              "summaries",
              f"lr_{_BASE_LR.value}_decay_{_WEIGHT_DECAY.value}",
              "clf_test",
          )
      )
  
      @tf.function
      def clf_train_step(iterator):
          def step_fn(inputs):
              clean_images = inputs["image_Clean"]
              all_obfs = tf.stack(
                  [
                      inputs["image_{}".format(obf)]
                      for obf in _OBFUSCATIONS_TRAIN.value
                  ],
                  axis=1,
              )
  
              # Randomly pick an obfuscation for each image - the first obfuscation is
              # the clean obfuscation.
              indices = tf.random.uniform(
                  shape=[all_obfs.shape[0]],
                  minval=1,
                  maxval=len(_OBFUSCATIONS_TRAIN.value),
                  dtype=tf.int32,
              )
              obf_images = tf.gather(all_obfs, indices, batch_dims=1)
              labels = inputs["label"]
  
              if clean_images.dtype == tf.uint8:
                  clean_images = (
                      tf.cast(clean_images, tf.float32) / 255.0
                  )
              if obf_images.dtype == tf.uint8:
                  obf_images = tf.cast(obf_images, tf.float32) / 255.0
  
              with tf.GradientTape() as tape:
                  clean_embeddings = (
                      model_clf.feature_extractor.encode_clean(
                          clean_images, training=True
                      )
                  )
                  if _EMBED_MAPPING_TYPE.value == "MLP":
                      obf_embeddings = model_clf.feature_extractor(
                          obf_images, training=True
                      )
                      # Embedding matching loss.
                      embed_loss = tf.reduce_mean(
                          tf.keras.metrics.mean_squared_error(
                              clean_embeddings, obf_embeddings
                          )
                      )
                  elif _EMBED_MAPPING_TYPE.value == "VAE":
                      (
                          obf_embeddings,
                          z_mean,
                          z_log_var,
                      ) = model_clf.feature_extractor(
                          obf_images, training=True
                      )
  
                      # Calculate appropriate KL term weight (based on a sigmoid, which
                      # ramps up to approximately 1 in _KL_LIM epochs). More specifically,
                      # this sigmoid is chosen so that it spans the range [exp(-5), exp(5)],
                      # approximately [0,1], within _KL_LIM epochs. See Figure 2 of
                      # https://arxiv.org/pdf/1511.06349.pdf.
                      kl_weight = tf.sigmoid(
                          10.0
                          * (
                              tf.cast(optimizer.iterations, tf.float32)
                              - 0.5
                              * tf.cast(
                                  _KL_LIM.value * train_steps_per_epoch,
                                  tf.float32,
                              )
                          )
                          / tf.cast(
                              _KL_LIM.value * train_steps_per_epoch,
                              tf.float32,
                          )
                      )
  
                      # Reconstruction loss and KL divergence loss for the VAE model.
                      recon_loss = tf.reduce_mean(
                          tf.keras.metrics.mean_squared_error(
                              clean_embeddings, obf_embeddings
                          )
                      )
                      kl_loss = losses_lib.vae_kl_divergence(
                          z_mean, z_log_var
                      )
                      embed_loss = recon_loss + kl_weight * kl_loss
  
                  # Crossentropy loss for the classifier.
                  logits = model_clf.clf_layer(
                      obf_embeddings, training=True
                  )
                  crossentropy_loss = tf.reduce_mean(
                      tf.nn.sparse_softmax_cross_entropy_with_logits(
                          labels, logits
                      )
                  )
  
                  # model_clf.losses contains the regularizer for the linear classifier.
                  regularizer_loss = tf.reduce_sum(model_clf.losses)
  
                  loss = (
                      _EMBED_LOSS_WEIGHT.value * embed_loss
                      + crossentropy_loss
                      + regularizer_loss
                  )
  
                  # Divide by number of replicas to balance out the artificially increased
                  # batch size.
                  per_replica_loss = (
                      loss / strategy.num_replicas_in_sync
                  )
  
              gradients = tape.gradient(
                  per_replica_loss, model_clf.trainable_variables
              )
              optimizer.apply_gradients(
                  zip(gradients, model_clf.trainable_variables)
              )
              loss_metrics["clf_train_loss"].update_state(loss)
              loss_metrics["clf_train_embed_loss"].update_state(
                  embed_loss
              )
              loss_metrics[
                  "clf_train_classification_loss"
              ].update_state(crossentropy_loss)
              loss_metrics["clf_train_reg_loss"].update_state(
                  regularizer_loss
              )
              acc_metrics["clf_train_acc"].update_state(
                  labels, tf.nn.softmax(logits)
              )
  
          strategy.run(step_fn, args=(next(iterator),))
  
      @tf.function
      def clf_test_step(item):
          def step_fn(inputs):
              labels = inputs["label"]
              for obf in _OBFUSCATIONS_EVAL.value:
                  images = inputs["image_{}".format(obf)]
                  if images.dtype == tf.uint8:
                      images = tf.cast(images, tf.float32) / 255.0
  
                  logits = model_clf(images, training=False)
  
                  losses = (
                      tf.nn.sparse_softmax_cross_entropy_with_logits(
                          labels, logits
                      )
                  )
                  loss = tf.math.divide_no_nan(
                      tf.reduce_sum(losses),
                      tf.cast(tf.size(losses), dtype=tf.float32),
                  )
  
                  loss_metrics[
                      "clf_test_loss_{}".format(obf)
                  ].update_state(loss)
                  acc_metrics[
                      "clf_test_acc_{}".format(obf)
                  ].update_state(labels, tf.nn.softmax(logits))
  
          strategy.run(step_fn, args=(item,))
  
      save_sample_images(
          train_dataset, clf_train_summary_writer, strategy
      )
      train_iterator = iter(train_dataset)
      for epoch in range(_EPOCHS.value):
          logging.info("Training Epoch: %s", epoch)
          with clf_train_summary_writer.as_default():
              for step in range(train_steps_per_epoch):
                  if step % _LOGGING_STEP.value == 0:
                      logging.info(
                          "Training Epoch %s, Step %s", epoch, step
                      )
                  clf_train_step(train_iterator)
  
              logging.info(
                  "Train Loss: %s, Train Accuracy %s %%",
                  round(
                      loss_metrics["clf_train_loss"].result().numpy(), 3
                  ),
                  round(
                      acc_metrics["clf_train_acc"].result().numpy()
                      * 100.0,
                      3,
                  ),
              )
  
              for loss_name, loss in loss_metrics.items():
                  if "train" in loss_name:
                      tf.summary.scalar(
                          loss_name,
                          loss.result().numpy(),
                          step=optimizer.iterations,
                      )
                      loss.reset_state()
  
              lr_sched = tf.keras.optimizers.schedules.ExponentialDecay.from_config(
                  optimizer.get_config()["learning_rate"]["config"]
              )
              tf.summary.scalar(
                  "clf_learning_rate",
                  lr_sched(optimizer.iterations).numpy(),
                  step=optimizer.iterations,
              )
  
              tf.summary.scalar(
                  "clf_train_acc",
                  acc_metrics["clf_train_acc"].result().numpy(),
                  step=optimizer.iterations,
              )
              acc_metrics["clf_train_acc"].reset_state()
  
          logging.info("Testing Epoch: %s", epoch)
          with clf_test_summary_writer.as_default():
              test_iterator = iter(test_dataset)
              for step in range(eval_steps_per_epoch):
                  item = next(test_iterator)
                  if step % _LOGGING_STEP.value == 0:
                      logging.info(
                          "Testing Epoch %s, Step %s", epoch, step
                      )
                  clf_test_step(item)
  
              for loss_name, loss in loss_metrics.items():
                  if "test" in loss_name:
                      tf.summary.scalar(
                          loss_name,
                          loss.result().numpy(),
                          step=optimizer.iterations,
                      )
                      loss.reset_state()
  
              for acc_name, acc in acc_metrics.items():
                  if "test" in acc_name:
                      tf.summary.scalar(
                          acc_name,
                          acc.result().numpy(),
                          step=optimizer.iterations,
                      )
                      acc.reset_state()
  
          clf_checkpoint.save(
              os.path.join(_MODEL_DIR.value, "clf_checkpoint")
          )
  
file_path: neural-structured-learning/research/meo/mlp_baseline/mlp_baseline.py
imports: import os
  from typing import Sequence
  from absl import app
  from absl import flags
  from absl import logging
  from meo.mlp_baseline import configs
  from meo.mlp_baseline import data_utils
  from meo.mlp_baseline import extended_model
  from meo.mlp_baseline import losses
  from meo.mlp_baseline import obfuscations
  import numpy
  import tensorflow
label_desc: Use image node_id 2 for calling a global function with example usage: train_and_eval_model(model_clf, train_dataset, train_steps_per_epoch, test_dataset, eval_steps_per_epoch, optimizer, strategy) without return types