class_docstr: Class with all code to run experiment.
class_docstr_tok: ["Class","with","all","code","to","run","experiment","."]
docstr: Runs training and evaluation.
docstr_tok: ["Runs","training","and","evaluation","."]
code: def train_and_eval(self):
      """Runs training and evaluation."""
      train_ds = imagenet_data.load(
          imagenet_data.Split.TRAIN,
          is_training=True,
          batch_dims=[
              jax.local_device_count() * FLAGS.train_device_batch_size
          ],
          tfds_data_dir=FLAGS.tfds_data_dir,
      )
  
      steps_per_epoch = (
          imagenet_data.Split.TRAIN.num_examples / self.total_batch_size
      )
      total_train_steps = int(steps_per_epoch * FLAGS.num_train_epochs)
      eval_every_n_steps = FLAGS.eval_every_n_steps
  
      checkpoint = objax.io.Checkpoint(FLAGS.model_dir, keep_ckpts=10)
      start_step, _ = checkpoint.restore(self.all_vars)
      cur_epoch = np.zeros([jax.local_device_count()], dtype=np.float32)
      for big_step in range(
          start_step, total_train_steps, eval_every_n_steps
      ):
          print(
              f"Running training steps {big_step + 1} - {big_step + eval_every_n_steps}"
          )
          with self.all_vars.replicate():
              # training
              start_time = time.time()
              for cur_step in range(
                  big_step + 1, big_step + eval_every_n_steps + 1
              ):
                  batch = next(train_ds)
                  cur_epoch[:] = cur_step / steps_per_epoch
                  monitors = self.train_op_parallel(
                      batch["images"], batch["labels"], cur_epoch
                  )
              elapsed_train_time = time.time() - start_time
              # eval
              start_time = time.time()
              accuracy = self.run_eval()
              elapsed_eval_time = time.time() - start_time
          # In multi-host setup only first host saves summaries and checkpoints.
          if jax.host_id() == 0:
              # save summary
              summary = objax.jaxboard.Summary()
              for k, v in monitors.items():
                  summary.scalar(f"train/{k}", v)
              # # Uncomment following two lines to save summary with training images
              # summary.image('input/train_img',
              #               imagenet_data.normalize_image_for_view(batch['images'][0]))
              summary.scalar("test/accuracy", accuracy * 100)
              self.summary_writer.write(summary, step=cur_step)
              # save checkpoint
              checkpoint.save(self.all_vars, cur_step)
          # print info
          print(
              "Step %d -- Epoch %.2f -- Loss %.2f Train Accuracy %.2f Eval Accuracy %.2f"
              % (
                  cur_step,
                  cur_step / steps_per_epoch,
                  monitors["total_loss"],
                  monitors["accuracy"],
                  accuracy * 100,
              )
          )
          print(
              "    Training took %.1f seconds, eval took %.1f seconds"
              % (elapsed_train_time, elapsed_eval_time),
              flush=True,
          )
  
file_path: objax/examples/image_classification/imagenet_resnet50_train.py
imports: import os
  import time
  import jax
  import jax.numpy
  import numpy
  from absl import app, flags, logging
  import objax
  from examples.image_classification import imagenet_resnet50_data
  from objax.zoo.resnet_v2 import ResNet50
label_desc: Use image node_id 7 for calling the Experiment obj's underlying member method code with example usage: obj.train_and_eval() without return types