code: def main():
      def batch_func(batch, ctx):
          data = gluon.utils.split_and_load(
              batch.data[0], ctx_list=ctx, batch_axis=0
          )
          label = gluon.utils.split_and_load(
              batch.label[0], ctx_list=ctx, batch_axis=0
          )
          return data, label
  
      opt = parse_args()
  
      filehandler = logging.FileHandler(opt.logging_file)
      streamhandler = logging.StreamHandler()
  
      logger = logging.getLogger("")
      logger.setLevel(logging.INFO)
      logger.addHandler(filehandler)
      logger.addHandler(streamhandler)
      if opt.data_backend[0:5] == "dali-":
          opt = parse_args(logger)  # Adding DALI parameters
  
      if opt.data_backend == "dali-gpu" and opt.num_gpus == 0:
          stream = os.popen("nvidia-smi -L | wc -l")
          opt.num_gpus = int(stream.read())
          logger.info(
              "When '--data-backend' is equal to 'dali-gpu', then `--num-gpus` should NOT be 0\n"
              "For now '--num-gpus' will be set to the number of GPUs installed: %d"
              % opt.num_gpus
          )
  
      logger.info(opt)
  
      classes = 1000
      num_gpus = opt.num_gpus
      batch_size = opt.batch_size * max(1, num_gpus)
      context = (
          [mx.gpu(i) for i in range(num_gpus)]
          if num_gpus > 0
          else [mx.cpu()]
      )
      opt.gpus = [i for i in range(num_gpus)]
  
      lr_decay = opt.lr_decay
      lr_decay_period = opt.lr_decay_period
      if opt.lr_decay_period > 0:
          lr_decay_epoch = list(
              range(lr_decay_period, opt.num_epochs, lr_decay_period)
          )
      else:
          lr_decay_epoch = [
              int(i) for i in opt.lr_decay_epoch.split(",")
          ]
      lr_decay_epoch = [e - opt.warmup_epochs for e in lr_decay_epoch]
      num_batches = opt.num_training_samples // batch_size
  
      lr_scheduler = LRSequential(
          [
              LRScheduler(
                  "linear",
                  base_lr=0,
                  target_lr=opt.lr,
                  nepochs=opt.warmup_epochs,
                  iters_per_epoch=num_batches,
              ),
              LRScheduler(
                  opt.lr_mode,
                  base_lr=opt.lr,
                  target_lr=0,
                  nepochs=opt.num_epochs - opt.warmup_epochs,
                  iters_per_epoch=num_batches,
                  step_epoch=lr_decay_epoch,
                  step_factor=lr_decay,
                  power=2,
              ),
          ]
      )
  
      model_name = opt.model
  
      kwargs = {
          "ctx": context,
          "pretrained": opt.use_pretrained,
          "classes": classes,
      }
      if opt.use_gn:
          kwargs["norm_layer"] = gcv.nn.GroupNorm
      if model_name.startswith("vgg"):
          kwargs["batch_norm"] = opt.batch_norm
      elif model_name.startswith("resnext"):
          kwargs["use_se"] = opt.use_se
  
      if opt.last_gamma:
          kwargs["last_gamma"] = True
  
      optimizer = "nag"
      optimizer_params = {
          "wd": opt.wd,
          "momentum": opt.momentum,
          "lr_scheduler": lr_scheduler,
      }
      if opt.dtype != "float32":
          optimizer_params["multi_precision"] = True
  
      net = get_model(model_name, **kwargs)
      net.cast(opt.dtype)
      if opt.resume_params != "":
          net.load_parameters(opt.resume_params, ctx=context)
  
      # teacher model for distillation training
      distillation = opt.teacher is not None and opt.hard_weight < 1.0
      if distillation:
          teacher = get_model(
              opt.teacher, pretrained=True, classes=classes, ctx=context
          )
          teacher.cast(opt.dtype)
  
      # Two functions for reading data from record file or raw images
      def get_data_rec(args):
          rec_train = os.path.expanduser(args.rec_train)
          rec_train_idx = os.path.expanduser(args.rec_train_idx)
          rec_val = os.path.expanduser(args.rec_val)
          rec_val_idx = os.path.expanduser(args.rec_val_idx)
          num_gpus = args.num_gpus
          batch_size = args.batch_size * max(1, num_gpus)
  
          jitter_param = 0.4
          lighting_param = 0.1
          input_size = opt.input_size
          crop_ratio = opt.crop_ratio if opt.crop_ratio > 0 else 0.875
          resize = int(math.ceil(input_size / crop_ratio))
  
          train_data = mx.io.ImageRecordIter(
              path_imgrec=rec_train,
              path_imgidx=rec_train_idx,
              preprocess_threads=args.num_workers,
              shuffle=True,
              batch_size=batch_size,
              data_shape=(3, input_size, input_size),
              mean_r=args.rgb_mean[0],
              mean_g=args.rgb_mean[1],
              mean_b=args.rgb_mean[2],
              std_r=args.rgb_std[0],
              std_g=args.rgb_std[1],
              std_b=args.rgb_std[2],
              rand_mirror=True,
              random_resized_crop=True,
              max_aspect_ratio=4.0 / 3.0,
              min_aspect_ratio=3.0 / 4.0,
              max_random_area=1,
              min_random_area=0.08,
              brightness=jitter_param,
              saturation=jitter_param,
              contrast=jitter_param,
              pca_noise=lighting_param,
          )
          val_data = mx.io.ImageRecordIter(
              path_imgrec=rec_val,
              path_imgidx=rec_val_idx,
              preprocess_threads=args.num_workers,
              shuffle=False,
              batch_size=batch_size,
              resize=resize,
              data_shape=(3, input_size, input_size),
              mean_r=args.rgb_mean[0],
              mean_g=args.rgb_mean[1],
              mean_b=args.rgb_mean[2],
              std_r=args.rgb_std[0],
              std_g=args.rgb_std[1],
              std_b=args.rgb_std[2],
          )
          return train_data, val_data, batch_func
  
      def get_data_rec_transfomed(args):
          data_dir = args.data_dir
          num_workers = args.num_workers
          batch_size = args.batch_size * max(1, args.num_gpus)
          normalize = transforms.Normalize(
              [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
          )
          jitter_param = 0.4
          lighting_param = 0.1
          input_size = opt.input_size
          crop_ratio = opt.crop_ratio if opt.crop_ratio > 0 else 0.875
          resize = int(math.ceil(input_size / crop_ratio))
  
          def batch_fn(batch, ctx):
              data = gluon.utils.split_and_load(
                  batch[0], ctx_list=ctx, batch_axis=0
              )
              label = gluon.utils.split_and_load(
                  batch[1], ctx_list=ctx, batch_axis=0
              )
              return data, label
  
          transform_train = transforms.Compose(
              [
                  transforms.RandomResizedCrop(input_size),
                  transforms.RandomFlipLeftRight(),
                  transforms.RandomColorJitter(
                      brightness=jitter_param,
                      contrast=jitter_param,
                      saturation=jitter_param,
                  ),
                  transforms.RandomLighting(lighting_param),
                  transforms.ToTensor(),
                  normalize,
              ]
          )
          transform_test = transforms.Compose(
              [
                  transforms.Resize(resize, keep_ratio=True),
                  transforms.CenterCrop(input_size),
                  transforms.ToTensor(),
                  normalize,
              ]
          )
  
          train_data = gluon.data.DataLoader(
              imagenet.classification.ImageNet(
                  data_dir, train=True
              ).transform_first(transform_train),
              batch_size=batch_size,
              shuffle=True,
              last_batch="discard",
              num_workers=num_workers,
          )
          val_data = gluon.data.DataLoader(
              imagenet.classification.ImageNet(
                  data_dir, train=False
              ).transform_first(transform_test),
              batch_size=batch_size,
              shuffle=False,
              num_workers=num_workers,
          )
  
          return train_data, val_data, batch_fn
  
      def get_data_loader(args):
          if args.data_backend == "mxnet":
              return (
                  get_data_rec
                  if args.use_rec
                  else get_data_rec_transfomed
              )
  
          # Check if DALI is installed:
          if args.data_backend[0:5] == "dali-":
              import dali
          if args.data_backend == "dali-gpu":
              return lambda *args, **kwargs: dali.get_rec_iter(
                  *args, **kwargs, batch_fn=batch_func, dali_cpu=False
              )
          if args.data_backend == "dali-cpu":
              return lambda *args, **kwargs: dali.get_rec_iter(
                  *args, **kwargs, batch_fn=batch_func, dali_cpu=True
              )
          raise ValueError("Wrong data backend")
  
      train_data, val_data, batch_fn = get_data_loader(opt)(opt)
      train_metric = (
          mx.metric.RMSE() if opt.mixup else mx.metric.Accuracy()
      )
      acc_top1 = mx.metric.Accuracy()
      acc_top5 = mx.metric.TopKAccuracy(5)
  
      save_frequency = opt.save_frequency
      if opt.save_dir and save_frequency:
          save_dir = opt.save_dir
          makedirs(save_dir)
      else:
          save_dir = ""
          save_frequency = 0
  
      def mixup_transform(label, classes, lam=1, eta=0.0):
          if isinstance(label, nd.NDArray):
              label = [label]
          res = []
          for l in label:
              y1 = l.one_hot(
                  classes,
                  on_value=1 - eta + eta / classes,
                  off_value=eta / classes,
              )
              y2 = l[::-1].one_hot(
                  classes,
                  on_value=1 - eta + eta / classes,
                  off_value=eta / classes,
              )
              res.append(lam * y1 + (1 - lam) * y2)
          return res
  
      def smooth(label, classes, eta=0.1):
          if isinstance(label, nd.NDArray):
              label = [label]
          smoothed = []
          for l in label:
              res = l.one_hot(
                  classes,
                  on_value=1 - eta + eta / classes,
                  off_value=eta / classes,
              )
              smoothed.append(res)
          return smoothed
  
      def test(ctx, val_data, val_batch):
          if opt.use_rec:
              val_data.reset()
          acc_top1.reset()
          acc_top5.reset()
          for i, batch in enumerate(val_data):
              for j in range(val_batch):
                  data, label = (
                      batch_fn(batch[j], ctx)
                      if type(batch) == list
                      else batch_fn(batch, ctx)
                  )
                  outputs = [
                      net(X.astype(opt.dtype, copy=False)) for X in data
                  ]
                  acc_top1.update(label, outputs)
                  acc_top5.update(label, outputs)
  
          _, top1 = acc_top1.get()
          _, top5 = acc_top5.get()
          return (1 - top1, 1 - top5)
  
      def train(ctx):
          if isinstance(ctx, mx.Context):
              ctx = [ctx]
          if opt.resume_params == "":
              net.initialize(mx.init.MSRAPrelu(), ctx=ctx)
  
          if opt.no_wd:
              for k, v in net.collect_params(
                  ".*beta|.*gamma|.*bias"
              ).items():
                  v.wd_mult = 0.0
  
          trainer = gluon.Trainer(
              net.collect_params(), optimizer, optimizer_params
          )
          if opt.resume_states != "":
              trainer.load_states(opt.resume_states)
  
          sparse_label_loss = not (opt.label_smoothing or opt.mixup)
          if distillation:
              L = gcv.loss.DistillationSoftmaxCrossEntropyLoss(
                  temperature=opt.temperature,
                  hard_weight=opt.hard_weight,
                  sparse_label=sparse_label_loss,
              )
          else:
              L = gluon.loss.SoftmaxCrossEntropyLoss(
                  sparse_label=sparse_label_loss
              )
  
          best_val_score = 1
          eta = 0.1 if opt.label_smoothing else 0.0
          start_time = time.time()
          val_batch = len(ctx) if opt.data_backend != "mxnet" else 1
          for epoch in range(opt.resume_epoch, opt.num_epochs):
              tic = time.time()
              if opt.use_rec:
                  train_data.reset()
              train_metric.reset()
              btic = time.time()
  
              for i, batch in enumerate(train_data):
                  for j in range(val_batch):
                      data, label = (
                          batch_fn(batch[j], ctx)
                          if type(batch) == list
                          else batch_fn(batch, ctx)
                      )
                      if opt.mixup:
                          lam = np.random.beta(
                              opt.mixup_alpha, opt.mixup_alpha
                          )
                          if (
                              epoch
                              >= opt.num_epochs - opt.mixup_off_epoch
                          ):
                              lam = 1
                          data = [
                              lam * X + (1 - lam) * X[::-1]
                              for X in data
                          ]
                          label = mixup_transform(
                              label, classes, lam, eta
                          )
  
                      elif opt.label_smoothing:
                          hard_label = label
                          label = smooth(label, classes)
  
                      if distillation:
                          teacher_prob = [
                              nd.softmax(
                                  teacher(
                                      X.astype(opt.dtype, copy=False)
                                  )
                                  / opt.temperature
                              )
                              for X in data
                          ]
  
                      with ag.record():
                          outputs = [
                              net(X.astype(opt.dtype, copy=False))
                              for X in data
                          ]
                          if distillation:
                              loss = [
                                  L(
                                      yhat.astype(
                                          "float32", copy=False
                                      ),
                                      y.astype("float32", copy=False),
                                      p.astype("float32", copy=False),
                                  )
                                  for yhat, y, p in zip(
                                      outputs, label, teacher_prob
                                  )
                              ]
                          else:
                              loss = [
                                  L(
                                      yhat,
                                      y.astype(opt.dtype, copy=False),
                                  )
                                  for yhat, y in zip(outputs, label)
                              ]
                      for l in loss:
                          l.backward()
                      trainer.step(batch_size)
  
                      if opt.mixup:
                          output_softmax = [
                              nd.SoftmaxActivation(
                                  out.astype("float32", copy=False)
                              )
                              for out in outputs
                          ]
                          train_metric.update(label, output_softmax)
                      else:
                          if opt.label_smoothing:
                              train_metric.update(hard_label, outputs)
                          else:
                              train_metric.update(label, outputs)
  
                  if (
                      opt.log_interval
                      and not (i + 1) % opt.log_interval
                  ):
                      (
                          train_metric_name,
                          train_metric_score,
                      ) = train_metric.get()
                      logger.info(
                          "Epoch[%d] Batch [%d]\tSpeed: %f samples/sec\t%s=%f\tlr=%f"
                          % (
                              epoch,
                              i + 1,
                              batch_size
                              * opt.log_interval
                              * val_batch
                              / (time.time() - btic),
                              train_metric_name,
                              train_metric_score,
                              trainer.learning_rate,
                          )
                      )
                      btic = time.time()
  
              train_metric_name, train_metric_score = train_metric.get()
              if opt.log_interval and i % opt.log_interval:
                  # We did NOT report the speed on the last iteration of the loop. Let's do it now
                  logger.info(
                      "Epoch[%d] Batch [%d]\tSpeed: %f samples/sec\t%s=%f\tlr=%f"
                      % (
                          epoch,
                          i,
                          batch_size
                          * (i % opt.log_interval)
                          * val_batch
                          / (time.time() - btic),
                          train_metric_name,
                          train_metric_score,
                          trainer.learning_rate,
                      )
                  )
  
              epoch_time = time.time() - tic
              throughput = int(batch_size * i * val_batch / epoch_time)
              err_top1_val, err_top5_val = test(
                  ctx, val_data, val_batch
              )
  
              logger.info(
                  "[Epoch %d] training: %s=%f"
                  % (epoch, train_metric_name, train_metric_score)
              )
              logger.info(
                  "[Epoch %d] speed: %d samples/sec\ttime cost: %f"
                  % (epoch, throughput, epoch_time)
              )
              logger.info(
                  "[Epoch %d] validation: err-top1=%f err-top5=%f"
                  % (epoch, err_top1_val, err_top5_val)
              )
  
              if err_top1_val < best_val_score:
                  best_val_score = err_top1_val
                  net.save_parameters(
                      "%s/%.4f-imagenet-%s-%d-best.params"
                      % (save_dir, best_val_score, model_name, epoch)
                  )
                  trainer.save_states(
                      "%s/%.4f-imagenet-%s-%d-best.states"
                      % (save_dir, best_val_score, model_name, epoch)
                  )
  
              if (
                  save_frequency
                  and save_dir
                  and (epoch + 1) % save_frequency == 0
              ):
                  net.save_parameters(
                      "%s/imagenet-%s-%d.params"
                      % (save_dir, model_name, epoch)
                  )
                  trainer.save_states(
                      "%s/imagenet-%s-%d.states"
                      % (save_dir, model_name, epoch)
                  )
  
          if save_frequency and save_dir:
              net.save_parameters(
                  "%s/imagenet-%s-%d.params"
                  % (save_dir, model_name, opt.num_epochs - 1)
              )
              trainer.save_states(
                  "%s/imagenet-%s-%d.states"
                  % (save_dir, model_name, opt.num_epochs - 1)
              )
  
          logger.info(
              "Training time for %d epochs: %f sec."
              % (opt.num_epochs, time.time() - start_time)
          )
  
      if opt.mode == "hybrid":
          net.hybridize(static_alloc=True, static_shape=True)
          if distillation:
              teacher.hybridize(static_alloc=True, static_shape=True)
      train(context)
  
file_path: gluon-cv/scripts/classification/imagenet/train_imagenet.py
imports: import argparse, time, logging, os, math
  import numpy
  import mxnet
  from mxnet import gluon, nd
  from mxnet import autograd
  from mxnet.gluon.data.vision import transforms
  import gluoncv
  from gluoncv.data import imagenet
  from gluoncv.model_zoo import get_model
  from gluoncv.utils import makedirs, LRSequential, LRScheduler
label_desc: Use image node_id 2 for calling a global function with example usage: main() and returns: data,  label, train_data,  val_data,  batch_func, train_data,  val_data,  batch_fn, res, smoothed, unknown,  unknown, data,  label, unknown, function, function