code: def main():
      global best_prec1, args
      best_prec1 = 0
      args = parse()
  
      if not len(args.data):
          raise Exception("error: No data set provided")
  
      if args.test:
          print("Test mode - only 10 iterations")
  
      args.distributed = False
      if "WORLD_SIZE" in os.environ:
          args.distributed = int(os.environ["WORLD_SIZE"]) > 1
      if "LOCAL_RANK" in os.environ:
          args.local_rank = int(os.environ["LOCAL_RANK"])
      else:
          args.local_rank = 0
  
      print("fp16_mode = {}".format(args.fp16_mode))
      print(
          "loss_scale = {}".format(args.loss_scale),
          type(args.loss_scale),
      )
  
      print(
          "\nCUDNN VERSION: {}\n".format(torch.backends.cudnn.version())
      )
  
      cudnn.benchmark = True
      best_prec1 = 0
      if args.deterministic:
          cudnn.benchmark = False
          cudnn.deterministic = True
          torch.manual_seed(args.local_rank)
          torch.set_printoptions(precision=10)
  
      args.gpu = 0
      args.world_size = 1
  
      if args.distributed:
          args.gpu = args.local_rank
          torch.cuda.set_device(args.gpu)
          torch.distributed.init_process_group(
              backend="nccl", init_method="env://"
          )
          args.world_size = torch.distributed.get_world_size()
  
      args.total_batch_size = args.world_size * args.batch_size
      assert (
          torch.backends.cudnn.enabled
      ), "Amp requires cudnn backend to be enabled."
  
      # create model
      if args.pretrained:
          print("=> using pre-trained model '{}'".format(args.arch))
          model = models.__dict__[args.arch](pretrained=True)
      else:
          print("=> creating model '{}'".format(args.arch))
          model = models.__dict__[args.arch]()
  
      if hasattr(torch, "channels_last") and hasattr(
          torch, "contiguous_format"
      ):
          if args.channels_last:
              memory_format = torch.channels_last
          else:
              memory_format = torch.contiguous_format
          model = model.cuda().to(memory_format=memory_format)
      else:
          model = model.cuda()
  
      # Scale learning rate based on global batch size
      args.lr = (
          args.lr * float(args.batch_size * args.world_size) / 256.0
      )
      optimizer = torch.optim.SGD(
          model.parameters(),
          args.lr,
          momentum=args.momentum,
          weight_decay=args.weight_decay,
      )
  
      if args.distributed:
          s = torch.cuda.Stream()
          s.wait_stream(torch.cuda.current_stream())
          with torch.cuda.stream(s):
              model = DDP(
                  model,
                  device_ids=[args.local_rank],
                  output_device=args.local_rank,
              )
          torch.cuda.current_stream().wait_stream(s)
  
      # define loss function (criterion) and optimizer
      criterion = nn.CrossEntropyLoss().cuda()
  
      # Optionally resume from a checkpoint
      if args.resume:
          # Use a local scope to avoid dangling references
          def resume():
              if os.path.isfile(args.resume):
                  print(
                      "=> loading checkpoint '{}'".format(args.resume)
                  )
                  checkpoint = torch.load(
                      args.resume,
                      map_location=lambda storage, loc: storage.cuda(
                          args.gpu
                      ),
                  )
                  args.start_epoch = checkpoint["epoch"]
                  global best_prec1
                  best_prec1 = checkpoint["best_prec1"]
                  model.load_state_dict(checkpoint["state_dict"])
                  optimizer.load_state_dict(checkpoint["optimizer"])
                  print(
                      "=> loaded checkpoint '{}' (epoch {})".format(
                          args.resume, checkpoint["epoch"]
                      )
                  )
              else:
                  print(
                      "=> no checkpoint found at '{}'".format(
                          args.resume
                      )
                  )
  
          resume()
  
      # Data loading code
      if len(args.data) == 1:
          traindir = os.path.join(args.data[0], "train")
          valdir = os.path.join(args.data[0], "val")
      else:
          traindir = args.data[0]
          valdir = args.data[1]
  
      if args.arch == "inception_v3":
          raise RuntimeError(
              "Currently, inception_v3 is not supported by this example."
          )
          # crop_size = 299
          # val_size = 320 # I chose this value arbitrarily, we can adjust.
      else:
          crop_size = 224
          val_size = 256
  
      train_loader = None
      val_loader = None
      if not args.disable_dali:
          train_pipe = create_dali_pipeline(
              batch_size=args.batch_size,
              num_threads=args.workers,
              device_id=args.local_rank,
              seed=12 + args.local_rank,
              data_dir=traindir,
              crop=crop_size,
              size=val_size,
              dali_cpu=args.dali_cpu,
              shard_id=args.local_rank,
              num_shards=args.world_size,
              is_training=True,
          )
          train_pipe.build()
          train_loader = DALIClassificationIterator(
              train_pipe,
              reader_name="Reader",
              last_batch_policy=LastBatchPolicy.PARTIAL,
              auto_reset=True,
          )
  
          val_pipe = create_dali_pipeline(
              batch_size=args.batch_size,
              num_threads=args.workers,
              device_id=args.local_rank,
              seed=12 + args.local_rank,
              data_dir=valdir,
              crop=crop_size,
              size=val_size,
              dali_cpu=args.dali_cpu,
              shard_id=args.local_rank,
              num_shards=args.world_size,
              is_training=False,
          )
          val_pipe.build()
          val_loader = DALIClassificationIterator(
              val_pipe,
              reader_name="Reader",
              last_batch_policy=LastBatchPolicy.PARTIAL,
              auto_reset=True,
          )
      else:
          train_dataset = datasets.ImageFolder(
              traindir,
              transforms.Compose(
                  [
                      transforms.RandomResizedCrop(crop_size),
                      transforms.RandomHorizontalFlip(),
                  ]
              ),
          )
          val_dataset = datasets.ImageFolder(
              valdir,
              transforms.Compose(
                  [
                      transforms.Resize(val_size),
                      transforms.CenterCrop(crop_size),
                  ]
              ),
          )
  
          train_sampler = None
          val_sampler = None
          if args.distributed:
              train_sampler = (
                  torch.utils.data.distributed.DistributedSampler(
                      train_dataset
                  )
              )
              val_sampler = (
                  torch.utils.data.distributed.DistributedSampler(
                      val_dataset
                  )
              )
  
          collate_fn = lambda b: fast_collate(b, memory_format)
  
          train_loader = torch.utils.data.DataLoader(
              train_dataset,
              batch_size=args.batch_size,
              shuffle=(train_sampler is None),
              num_workers=args.workers,
              pin_memory=True,
              sampler=train_sampler,
              collate_fn=collate_fn,
          )
  
          val_loader = torch.utils.data.DataLoader(
              val_dataset,
              batch_size=args.batch_size,
              shuffle=False,
              num_workers=args.workers,
              pin_memory=True,
              sampler=val_sampler,
              collate_fn=collate_fn,
          )
  
      if args.evaluate:
          validate(val_loader, model, criterion)
          return
  
      scaler = torch.cuda.amp.GradScaler(
          init_scale=args.loss_scale,
          growth_factor=2,
          backoff_factor=0.5,
          growth_interval=100,
          enabled=args.fp16_mode,
      )
      total_time = AverageMeter()
      for epoch in range(args.start_epoch, args.epochs):
          # train for one epoch
          avg_train_time = train(
              train_loader, model, criterion, scaler, optimizer, epoch
          )
          total_time.update(avg_train_time)
          if args.test:
              break
  
          # evaluate on validation set
          [prec1, prec5] = validate(val_loader, model, criterion)
  
          # remember best prec@1 and save checkpoint
          if args.local_rank == 0:
              is_best = prec1 > best_prec1
              best_prec1 = max(prec1, best_prec1)
              save_checkpoint(
                  {
                      "epoch": epoch + 1,
                      "arch": args.arch,
                      "state_dict": model.state_dict(),
                      "best_prec1": best_prec1,
                      "optimizer": optimizer.state_dict(),
                  },
                  is_best,
              )
              if epoch == args.epochs - 1:
                  print(
                      "##Top-1 {0}\n"
                      "##Top-5 {1}\n"
                      "##Perf  {2}".format(
                          prec1,
                          prec5,
                          args.total_batch_size / total_time.avg,
                      )
                  )
  
imports: import argparse
  import os
  import shutil
  import time
  import math
  import torch
  import torch.nn
  import torch.nn.parallel
  import torch.backends.cudnn
  import torch.distributed
  import torch.optim
  import torch.utils.data
  import torch.utils.data.distributed
  import numpy
  from torch.nn.parallel import DistributedDataParallel
  import torchvision.transforms
  import torchvision.datasets
  import torchvision.models
label_desc: Use image node_id 5 for calling a global function with example usage: main() without return types