code: def main(argv):
      del argv
      setup_tf()
      unlabeled_name = FLAGS.unlabeled or FLAGS.dataset.split(".")[0]
      unlabeled = DATASETS_UNLABELED()[unlabeled_name]()
      dataset = DATASETS_LABELED()[FLAGS.dataset]()
      testsets = [dataset.test]
      if FLAGS.test_extra:
          testsets += [
              DATASETS_LABELED()[x]().test
              for x in FLAGS.test_extra.split(",")
          ]
      module = FixMatch(
          dataset.nclass,
          network(FLAGS.arch),
          lr=FLAGS.lr,
          wd=FLAGS.wd,
          arch=FLAGS.arch,
          batch=FLAGS.batch,
          wu=FLAGS.wu,
          confidence=FLAGS.confidence,
          uratio=FLAGS.uratio,
          filters=FLAGS.filters,
          filters_max=FLAGS.filters_max,
          repeat=FLAGS.repeat,
          scales=FLAGS.scales
          or objax.util.ilog2(dataset.train.image_shape[0]) - 2,
      )
      logdir = (
          f"{FLAGS.dataset}/{unlabeled_name}/{FLAGS.augment}/{module.__class__.__name__}/%s"
          % (
              "_".join(
                  sorted("%s%s" % k for k in module.params.items())
              )
          )
      )
      logdir = os.path.join(FLAGS.logdir, logdir)
      valid = (
          dataset.valid.parse().batch(FLAGS.batch).nchw().prefetch(16)
      )
      test = {}
      for x in testsets:
          test.update(
              (k, v.parse().batch(FLAGS.batch).nchw().prefetch(16))
              for k, v in x.items()
          )
      labeled = dataset.train.repeat().shuffle(FLAGS.shuffle).parse()
      unlabeled = (
          unlabeled.train.repeat().shuffle(FLAGS.shuffle).parse()
      )
      augment_pool = get_augment(dataset.train)
      train = augment_pool(
          labeled,
          unlabeled,
          dataset.nclass,
          FLAGS.batch,
          FLAGS.uratio,
          predict=module.eval_op,
      )
      module.train(
          FLAGS.train_kimg,
          FLAGS.report_kimg,
          train,
          valid,
          test,
          logdir,
          FLAGS.keep_ckpts,
      )
  
file_path: objax/examples/fixmatch/fixmatch.py
imports: import os
  from typing import Callable
  import jax.numpy
  from absl import app, flags
  import objax
  from examples.fixmatch.libml.augment.augment import get_augment
  from examples.fixmatch.libml.data.fsl import DATASETS_LABELED
  from examples.fixmatch.libml.data.ssl import DATASETS_UNLABELED
  from examples.fixmatch.libml.models import ARCHS, network
  from examples.fixmatch.libml.train import TrainLoopSSL
  from examples.fixmatch.libml.util import setup_tf
label_desc: Use image node_id 1 for calling a global function with example usage: main(argv) without return types