class_docstr: Class with all code to run experiment.
class_docstr_tok: ["Class","with","all","code","to","run","experiment","."]
code: def __init__(self):
      # Some constants
      self.total_batch_size = (
          FLAGS.train_device_batch_size * jax.device_count()
      )
      self.base_learning_rate = (
          FLAGS.base_learning_rate * self.total_batch_size / 256
      )
      # Create model
      bn_cls = (
          objax.nn.SyncedBatchNorm2D
          if FLAGS.use_sync_bn
          else objax.nn.BatchNorm2D
      )
      self.model = ResNet50(
          in_channels=3,
          num_classes=NUM_CLASSES,
          normalization_fn=bn_cls,
      )
      self.model_vars = self.model.vars()
      print(self.model_vars)
      # Create parallel eval op
      self.evaluate_batch_parallel = objax.Parallel(
          self.evaluate_batch,
          self.model_vars,
          reduce=lambda x: x.sum(0),
      )
      # Create parallel training op
      self.optimizer = objax.optimizer.Momentum(
          self.model_vars, momentum=0.9, nesterov=True
      )
      self.compute_grads_loss = objax.GradValues(
          self.loss_fn, self.model_vars
      )
      self.all_vars = self.model_vars + self.optimizer.vars()
      self.train_op_parallel = objax.Parallel(
          self.train_op, self.all_vars, reduce=lambda x: x[0]
      )
      # Summary writer
      self.summary_writer = objax.jaxboard.SummaryWriter(
          os.path.join(FLAGS.model_dir, "tb")
      )
  
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 1 to create a new Experiment object  with example: obj = Experiment()