code: def main():
      global best_prec1, args
  
      args = parse()
      print("opt_level = {}".format(args.opt_level))
      print(
          "keep_batchnorm_fp32 = {}".format(args.keep_batchnorm_fp32),
          type(args.keep_batchnorm_fp32),
      )
      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.distributed = False
      if "WORLD_SIZE" in os.environ:
          args.distributed = int(os.environ["WORLD_SIZE"]) > 1
  
      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()
  
      assert (
          torch.backends.cudnn.enabled
      ), "Amp requires cudnn backend to be enabled."
  
      if args.channels_last:
          memory_format = torch.channels_last
      else:
          memory_format = torch.contiguous_format
  
      # 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 args.sync_bn:
          import apex
  
          print("using apex synced BN")
          model = apex.parallel.convert_syncbn_model(model)
  
      model = model.cuda().to(memory_format=memory_format)
  
      # 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,
      )
  
      # Initialize Amp.  Amp accepts either values or strings for the optional override arguments,
      # for convenient interoperation with argparse.
      model, optimizer = amp.initialize(
          model,
          optimizer,
          opt_level=args.opt_level,
          keep_batchnorm_fp32=args.keep_batchnorm_fp32,
          loss_scale=args.loss_scale,
      )
  
      # For distributed training, wrap the model with apex.parallel.DistributedDataParallel.
      # This must be done AFTER the call to amp.initialize.  If model = DDP(model) is called
      # before model, ... = amp.initialize(model, ...), the call to amp.initialize may alter
      # the types of model's parameters in a way that disrupts or destroys DDP's allreduce hooks.
      if args.distributed:
          # By default, apex.parallel.DistributedDataParallel overlaps communication with
          # computation in the backward pass.
          # model = DDP(model)
          # delay_allreduce delays all communication to the end of the backward pass.
          model = DDP(model, delay_allreduce=True)
  
      # 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
      traindir = os.path.join(args.data, "train")
      valdir = os.path.join(args.data, "val")
  
      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_dataset = datasets.ImageFolder(
          traindir,
          transforms.Compose(
              [
                  transforms.RandomResizedCrop(crop_size),
                  transforms.RandomHorizontalFlip(),
                  # transforms.ToTensor(), Too slow
                  # normalize,
              ]
          ),
      )
      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
  
      for epoch in range(args.start_epoch, args.epochs):
          if args.distributed:
              train_sampler.set_epoch(epoch)
  
          # train for one epoch
          train(train_loader, model, criterion, optimizer, epoch)
  
          # evaluate on validation set
          prec1 = 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,
              )
  
imports: import argparse
  import os
  import shutil
  import time
  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 torchvision.transforms
  import torchvision.datasets
  import torchvision.models
  import numpy
label_desc: Use image node_id 3 for calling a global function with example usage: main() without return types