code: def main(_):
      keras.mixed_precision.set_dtype_policy(
          FLAGS.mixed_precision_policy
      )
  
      logging.info(
          "Benchmarking configs...\n"
          "=========================\n"
          f"MODEL: {FLAGS.model}\n"
          f"TASK: image classification/dogs-vs-cats \n"
          f"BATCH_SIZE: {FLAGS.batch_size}\n"
          f"EPOCHS: {FLAGS.epochs}\n"
          "=========================\n"
      )
  
      # Load datasets.
      train_ds, validation_ds = load_data()
  
      # Load the model.
      classifier = load_model()
  
      lr = keras.optimizers.schedules.PolynomialDecay(
          5e-4,
          decay_steps=train_ds.cardinality() * FLAGS.epochs,
          end_learning_rate=0.0,
      )
      optimizer = keras.optimizers.Adam(lr)
      loss = keras.losses.SparseCategoricalCrossentropy(
          from_logits=True
      )
  
      benchmark_metrics_callback = BenchmarkMetricsCallback(
          start_batch=1,
          stop_batch=train_ds.cardinality().numpy() - 1,
      )
  
      classifier.compile(
          optimizer=optimizer,
          loss=loss,
          metrics=["sparse_categorical_accuracy"],
      )
      # Start training.
      logging.info("Starting Training...")
  
      st = time.time()
  
      history = classifier.fit(
          train_ds,
          validation_data=validation_ds,
          epochs=FLAGS.epochs,
          callbacks=[benchmark_metrics_callback],
      )
  
      wall_time = time.time() - st
      validation_accuracy = history.history[
          "val_sparse_categorical_accuracy"
      ][-1]
  
      examples_per_second = (
          np.mean(
              np.array(benchmark_metrics_callback.state["throughput"])
          )
          * FLAGS.batch_size
      )
  
      logging.info("Training Finished!")
      logging.info(f"Wall Time: {wall_time:.4f} seconds.")
      logging.info(f"Validation Accuracy: {validation_accuracy:.4f}")
      logging.info(f"examples_per_second: {examples_per_second:.4f}")
  
file_path: keras/benchmarks/model_benchmark/image_classification_benchmark.py
imports: import time
  import numpy
  import tensorflow
  import tensorflow_datasets
  from absl import app
  from absl import flags
  from absl import logging
  from model_benchmark.benchmark_utils import BenchmarkMetricsCallback
  import keras
label_desc: Use image node_id 3 for calling a global function with example usage: main(_) without return types