class_docstr: Training loop for general machine learning models.
  Based on the training loop from the objax CIFAR10 example code.
class_docstr_tok: ["Training","loop","for","general","machine","learning","models",".","Based","on","the","training","loop","from","the","objax","CIFAR10","example","code","."]
docstr: Completely standard training. Nothing interesting to see here.
docstr_tok: ["Completely","standard","training",".","Nothing","interesting","to","see","here","."]
code: def train(
      self,
      num_train_epochs: int,
      train_size: int,
      train: DataSet,
      test: DataSet,
      logdir: str,
      save_steps=100,
      patience=None,
  ):
      """
      Completely standard training. Nothing interesting to see here.
      """
      checkpoint = objax.io.Checkpoint(
          logdir, keep_ckpts=20, makedir=True
      )
      start_epoch, last_ckpt = checkpoint.restore(self.vars())
      train_iter = iter(train)
      progress = np.zeros(
          jax.local_device_count(), "f"
      )  # for multi-GPU
  
      best_acc = 0
      best_acc_epoch = -1
  
      with SummaryWriter(os.path.join(logdir, "tb")) as tensorboard:
          for epoch in range(start_epoch, num_train_epochs):
              # Train
              summary = Summary()
              loop = range(0, train_size, self.params.batch)
              for step in loop:
                  progress[:] = (step + (epoch * train_size)) / (
                      num_train_epochs * train_size
                  )
                  self.train_step(summary, next(train_iter), progress)
  
              # Eval
              accuracy, total = 0, 0
              if epoch % FLAGS.eval_steps == 0 and test is not None:
                  for data in test:
                      total += data["image"].shape[0]
                      preds = np.argmax(
                          self.predict(data["image"].numpy()), axis=1
                      )
                      accuracy += (preds == data["label"].numpy()).sum()
                  accuracy /= total
                  summary.scalar("eval/accuracy", 100 * accuracy)
                  tensorboard.write(
                      summary, step=(epoch + 1) * train_size
                  )
                  print(
                      "Epoch %04d  Loss %.2f  Accuracy %.2f"
                      % (
                          epoch + 1,
                          summary["losses/xe"](),
                          summary["eval/accuracy"](),
                      )
                  )
  
                  if summary["eval/accuracy"]() > best_acc:
                      best_acc = summary["eval/accuracy"]()
                      best_acc_epoch = epoch
                  elif (
                      patience is not None
                      and epoch > best_acc_epoch + patience
                  ):
                      print("early stopping!")
                      checkpoint.save(self.vars(), epoch + 1)
                      return
  
              else:
                  print(
                      "Epoch %04d  Loss %.2f  Accuracy --"
                      % (epoch + 1, summary["losses/xe"]())
                  )
  
              if epoch % save_steps == save_steps - 1:
                  checkpoint.save(self.vars(), epoch + 1)
  
imports: import functools
  import os
  import shutil
  from typing import Callable
  import json
  import jax
  import jax.numpy
  import numpy
  import tensorflow
  import tensorflow_datasets
  from absl import app, flags
  import objax
  from objax.jaxboard import SummaryWriter, Summary
  from objax.util import EasyDict
  from objax.zoo import convnet, wide_resnet
  from dataset import DataSet
label_desc: Use image node_id 3 for calling the TrainLoop obj's underlying member method code with example usage: obj.train(num_train_epochs, train_size, train, test, logdir, save_steps, patience) without return types