code: def main():
      opt = parse_args()
  
      bps.init()
  
      gpu_name = subprocess.check_output(
          ["nvidia-smi", "--query-gpu=gpu_name", "--format=csv"]
      )
      gpu_name = gpu_name.decode("utf8").split("\n")[-2]
      gpu_name = "-".join(gpu_name.split())
      filename = "cifar100-%d-%s-%s.log" % (
          bps.size(),
          gpu_name,
          opt.logging_file,
      )
      filehandler = logging.FileHandler(filename)
      streamhandler = logging.StreamHandler()
  
      logger = logging.getLogger("")
      logger.setLevel(logging.INFO)
      logger.addHandler(filehandler)
      logger.addHandler(streamhandler)
  
      logger.info(opt)
  
      batch_size = opt.batch_size
      classes = 100
  
      num_gpus = opt.num_gpus
      # batch_size *= max(1, num_gpus)
      context = (
          mx.gpu(bps.local_rank())
          if num_gpus > 0
          else mx.cpu(bps.local_rank())
      )
      num_workers = opt.num_workers
      nworker = bps.size()
      rank = bps.rank()
  
      lr_decay = opt.lr_decay
      lr_decay_epoch = [int(i) for i in opt.lr_decay_epoch.split(",")]
  
      num_batches = 50000 // (opt.batch_size * nworker)
      lr_scheduler = LRSequential(
          [
              LRScheduler(
                  "linear",
                  base_lr=opt.warmup_lr,
                  target_lr=opt.lr * nworker / bps.local_size(),
                  nepochs=opt.warmup_epochs,
                  iters_per_epoch=num_batches,
              ),
              LRScheduler(
                  "step",
                  base_lr=opt.lr * nworker / bps.local_size(),
                  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,
              ),
          ]
      )
  
      num_batches = 50000 // (opt.batch_size * nworker)
      lr_scheduler = LRSequential(
          [
              LRScheduler(
                  "linear",
                  base_lr=opt.warmup_lr,
                  target_lr=opt.lr * nworker / bps.local_size(),
                  nepochs=opt.warmup_epochs,
                  iters_per_epoch=num_batches,
              ),
              LRScheduler(
                  "step",
                  base_lr=opt.lr * nworker / bps.local_size(),
                  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
      if model_name.startswith("cifar_wideresnet"):
          kwargs = {"classes": classes, "drop_rate": opt.drop_rate}
      else:
          kwargs = {"classes": classes}
      net = get_model(model_name, **kwargs)
      if opt.resume_from:
          net.load_parameters(opt.resume_from, ctx=context)
  
      if opt.compressor:
          optimizer = "sgd"
      else:
          optimizer = "nag"
  
      save_period = opt.save_period
      if opt.save_dir and save_period:
          save_dir = opt.save_dir
          makedirs(save_dir)
      else:
          save_dir = ""
          save_period = 0
  
      # from https://github.com/weiaicunzai/pytorch-cifar/blob/master/conf/global_settings.py
      CIFAR100_TRAIN_MEAN = [
          0.5070751592371323,
          0.48654887331495095,
          0.4409178433670343,
      ]
      CIFAR100_TRAIN_STD = [
          0.2673342858792401,
          0.2564384629170883,
          0.27615047132568404,
      ]
  
      transform_train = transforms.Compose(
          [
              gcv_transforms.RandomCrop(32, pad=4),
              transforms.RandomFlipLeftRight(),
              transforms.ToTensor(),
              transforms.Normalize(
                  CIFAR100_TRAIN_MEAN, CIFAR100_TRAIN_STD
              ),
          ]
      )
  
      transform_test = transforms.Compose(
          [
              transforms.ToTensor(),
              transforms.Normalize(
                  CIFAR100_TRAIN_MEAN, CIFAR100_TRAIN_STD
              ),
          ]
      )
  
      def test(ctx, val_data):
          metric = mx.metric.Accuracy()
          for i, batch in enumerate(val_data):
              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
              )
              outputs = [net(X) for X in data]
              metric.update(label, outputs)
          return metric.get()
  
      def train(epochs, ctx):
          if isinstance(ctx, mx.Context):
              ctx = [ctx]
          net.initialize(mx.init.Xavier(), ctx=ctx)
  
          train_data = gluon.data.DataLoader(
              gluon.data.vision.CIFAR100(train=True)
              .shard(nworker, rank)
              .transform_first(transform_train),
              batch_size=batch_size,
              shuffle=True,
              last_batch="discard",
              num_workers=num_workers,
          )
  
          val_data = gluon.data.DataLoader(
              gluon.data.vision.CIFAR100(train=False)
              .shard(nworker, rank)
              .transform_first(transform_test),
              batch_size=batch_size,
              shuffle=False,
              num_workers=num_workers,
          )
  
          params = net.collect_params()
  
          compression_params = {
              "compressor": opt.compressor,
              "ef": opt.ef,
              "momentum": opt.compress_momentum,
              "scaling": opt.onebit_scaling,
              "k": opt.k,
              "fp16": opt.fp16_pushpull,
          }
  
          optimizer_params = {
              "lr_scheduler": lr_scheduler,
              "wd": opt.wd,
              "momentum": opt.momentum,
          }
  
          trainer = bps.DistributedTrainer(
              params,
              optimizer,
              optimizer_params,
              compression_params=compression_params,
          )
          metric = mx.metric.Accuracy()
          train_metric = mx.metric.Accuracy()
          loss_fn = gluon.loss.SoftmaxCrossEntropyLoss()
  
          iteration = 0
          best_val_score = 0
          bps.byteps_declare_tensor("acc")
          for epoch in range(epochs):
              tic = time.time()
              train_metric.reset()
              metric.reset()
              train_loss = 0
              num_batch = len(train_data)
  
              for i, batch in enumerate(train_data):
                  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
                  )
  
                  with ag.record():
                      output = [net(X) for X in data]
                      loss = [
                          loss_fn(yhat, y)
                          for yhat, y in zip(output, label)
                      ]
                  for l in loss:
                      l.backward()
                  trainer.step(batch_size)
                  train_loss += sum([l.sum().asscalar() for l in loss])
  
                  train_metric.update(label, output)
                  name, train_acc = train_metric.get()
                  iteration += 1
  
              train_loss /= batch_size * num_batch
              name, train_acc = train_metric.get()
              throughput = int(
                  batch_size * nworker * i / (time.time() - tic)
              )
  
              logger.info(
                  "[Epoch %d] speed: %d samples/sec\ttime cost: %f lr=%f"
                  % (
                      epoch,
                      throughput,
                      time.time() - tic,
                      trainer.learning_rate,
                  )
              )
  
              name, val_acc = test(ctx, val_data)
              acc = mx.nd.array([train_acc, val_acc], ctx=ctx[0])
              bps.byteps_push_pull(acc, name="acc", is_average=False)
              acc /= bps.size()
              train_acc, val_acc = acc[0].asscalar(), acc[1].asscalar()
              if bps.rank() == 0:
                  logger.info(
                      "[Epoch %d] training: %s=%f"
                      % (epoch, name, train_acc)
                  )
                  logger.info(
                      "[Epoch %d] validation: %s=%f"
                      % (epoch, name, val_acc)
                  )
  
              if val_acc > best_val_score:
                  best_val_score = val_acc
                  net.save_parameters(
                      "%s/%.4f-cifar-%s-%d-best.params"
                      % (save_dir, best_val_score, model_name, epoch)
                  )
  
              if (
                  save_period
                  and save_dir
                  and (epoch + 1) % save_period == 0
              ):
                  net.save_parameters(
                      "%s/cifar100-%s-%d.params"
                      % (save_dir, model_name, epoch)
                  )
  
          if save_period and save_dir:
              net.save_parameters(
                  "%s/cifar100-%s-%d.params"
                  % (save_dir, model_name, epochs - 1)
              )
  
      if opt.mode == "hybrid":
          net.hybridize()
      train(opt.num_epochs, context)
  
file_path: byteps/example/mxnet/train_cifar100_byteps_gc.py
imports: import argparse
  import logging
  import subprocess
  import time
  import gluoncv
  import matplotlib
  import mxnet
  from gluoncv.data import transforms
  from gluoncv.model_zoo import get_model
  from gluoncv.utils import LRScheduler, LRSequential, makedirs
  from mxnet import autograd
  from mxnet import gluon
  from mxnet.gluon.data.vision import transforms
  import byteps.mxnet
label_desc: Use image node_id 2 for calling a global function with example usage: main() and returns: metric