code: import argparse
  import os
  import shutil
  import time
  import math
  import logging
  
  import torch
  import torch.nn as nn
  import torch.nn.parallel
  import torch.backends.cudnn as cudnn
  import torch.distributed as dist
  import torch.optim
  import torch.utils.data
  import torch.utils.data.distributed
  import torchvision.transforms as transforms
  import torchvision.datasets as datasets
  import torchvision.models as models
  import numpy as np
  
  try:
      from apex.parallel import DistributedDataParallel as DDP
      from apex.fp16_utils import *
      from apex import amp, optimizers
      from apex.multi_tensor_apply import multi_tensor_applier
  except ImportError:
      raise ImportError("Please install apex from https://www.github.com/nvidia/apex to run this example.")
  
  
  try:
      from nvidia.dali.plugin.pytorch import DALIClassificationIterator
      from nvidia.dali.pipeline import Pipeline
      import nvidia.dali.ops as ops
      import nvidia.dali.types as types
  except ImportError:
      raise ImportError("Please install DALI from https://www.github.com/NVIDIA/DALI to run this example.")
  
  from nasws.cnn.utils import AverageMeter
  from utils import accuracy
  
  
  # item() is a recent addition, so this helps with backward compatibility.
  def to_python_float(t):
      if hasattr(t, 'item'):
          return t.item()
      else:
          return t[0]
  
  
  class HybridTrainPipe(Pipeline):
      def __init__(self, batch_size, num_threads, device_id, data_dir, crop,
                   shard_id, num_shards, dali_cpu=False, args=None,
                   file_list=None
                   ):
          
          super(HybridTrainPipe, self).__init__(batch_size,
                                                num_threads,
                                                device_id,
                                                seed=12 + device_id)
          
          self.input = ops.FileReader(file_root=data_dir,
                                      shard_id=args.apex_local_rank,
                                      num_shards=args.world_size,
                                      random_shuffle=True,
                                      pad_last_batch=True,
                                      file_list=file_list)
          #let user decide which pipeline works him bets for RN version he runs
          dali_device = 'cpu' if dali_cpu else 'gpu'
          decoder_device = 'cpu' if dali_cpu else 'mixed'
          # This padding sets the size of the internal nvJPEG buffers to be able to handle all images from full-sized ImageNet
          # without additional reallocations
          device_memory_padding = 211025920 if decoder_device == 'mixed' else 0
          host_memory_padding = 140544512 if decoder_device == 'mixed' else 0
          self.decode = ops.ImageDecoderRandomCrop(device=decoder_device, output_type=types.RGB,
                                                   device_memory_padding=device_memory_padding,
                                                   host_memory_padding=host_memory_padding,
                                                   random_aspect_ratio=[0.8, 1.25],
                                                   random_area=[0.1, 1.0],
                                                   num_attempts=100)
          self.res = ops.Resize(device=dali_device,
                                resize_x=crop,
                                resize_y=crop,
                                interp_type=types.INTERP_TRIANGULAR)
          self.cmnp = ops.CropMirrorNormalize(device="gpu",
                                              output_dtype=types.FLOAT,
                                              output_layout=types.NCHW,
                                              crop=(crop, crop),
                                              image_type=types.RGB,
                                              mean=[0.485 * 255,0.456 * 255,0.406 * 255],
                                              std=[0.229 * 255,0.224 * 255,0.225 * 255])
          self.coin = ops.CoinFlip(probability=0.5)
          logging.info('DALI "{0}" variant'.format(dali_device))
  
      def define_graph(self):
          rng = self.coin()
          self.jpegs, self.labels = self.input(name="Reader")
          images = self.decode(self.jpegs)
          images = self.res(images)
          output = self.cmnp(images.gpu(), mirror=rng)
          return [output, self.labels]
  
  
  class HybridValPipe(Pipeline):
      def __init__(self, batch_size, num_threads, device_id, data_dir, crop,
                   size, shard_id, num_shards, args=None):
          super(HybridValPipe, self).__init__(batch_size,
                                             num_threads,
                                              device_id,
                                              seed=12 + device_id)
          self.input = ops.FileReader(file_root=data_dir,
                                      shard_id=args.apex_local_rank,
                                      num_shards=args.world_size,
                                      random_shuffle=False,
                                      pad_last_batch=True)
          self.decode = ops.ImageDecoder(device="mixed", output_type=types.RGB)
          self.res = ops.Resize(device="gpu",
                                resize_shorter=size,
                                interp_type=types.INTERP_TRIANGULAR)
          self.cmnp = ops.CropMirrorNormalize(device="gpu",
                                              output_dtype=types.FLOAT,
                                              output_layout=types.NCHW,
                                              crop=(crop, crop),
                                              image_type=types.RGB,
                                              mean=[0.485 * 255,0.456 * 255,0.406 * 255],
                                              std=[0.229 * 255,0.224 * 255,0.225 * 255])
  
      def define_graph(self):
          self.jpegs, self.labels = self.input(name="Reader")
          images = self.decode(self.jpegs)
          images = self.res(images)
          output = self.cmnp(images)
          return [output, self.labels]
  
  
  def fast_collate(batch, memory_format):
  
      imgs = [img[0] for img in batch]
      targets = torch.tensor([target[1] for target in batch], dtype=torch.int64)
      w = imgs[0].size()[1]
      h = imgs[0].size()[2]
      # print(imgs[0].size())
      tensor = torch.zeros( (len(imgs), 3, h, w), dtype=torch.uint8).contiguous(memory_format=memory_format)
      for i, img in enumerate(imgs):
          nump_array = np.asarray(img, dtype=np.uint8)
          if(nump_array.ndim < 3):
              nump_array = np.expand_dims(nump_array, axis=-1)
          # nump_array = np.rollaxis(nump_array, 2)
          # print(nump_array.shape)
          tensor[i] += torch.from_numpy(nump_array)
      return tensor, targets
  
  
  class data_prefetcher():
      def __init__(self, loader):
          self.loader = iter(loader)
          self.stream = torch.cuda.Stream()
          self.mean = torch.tensor([0.485 * 255, 0.456 * 255, 0.406 * 255]).cuda().view(1,3,1,1)
          self.std = torch.tensor([0.229 * 255, 0.224 * 255, 0.225 * 255]).cuda().view(1,3,1,1)
          # With Amp, it isn't necessary to manually convert data to half.
          # if args.fp16:
          #     self.mean = self.mean.half()
          #     self.std = self.std.half()
          self.preload()
  
      def preload(self):
          try:
              self.next_input, self.next_target = next(self.loader)
          except StopIteration:
              self.next_input = None
              self.next_target = None
              return
          # if record_stream() doesn't work, another option is to make sure device inputs are created
          # on the main stream.
          # self.next_input_gpu = torch.empty_like(self.next_input, device='cuda')
          # self.next_target_gpu = torch.empty_like(self.next_target, device='cuda')
          # Need to make sure the memory allocated for next_* is not still in use by the main stream
          # at the time we start copying to next_*:
          # self.stream.wait_stream(torch.cuda.current_stream())
          with torch.cuda.stream(self.stream):
              self.next_input = self.next_input.cuda(non_blocking=True)
              self.next_target = self.next_target.cuda(non_blocking=True)
              # more code for the alternative if record_stream() doesn't work:
              # copy_ will record the use of the pinned source tensor in this side stream.
              # self.next_input_gpu.copy_(self.next_input, non_blocking=True)
              # self.next_target_gpu.copy_(self.next_target, non_blocking=True)
              # self.next_input = self.next_input_gpu
              # self.next_target = self.next_target_gpu
  
              # With Amp, it isn't necessary to manually convert data to half.
              # if args.fp16:
              #     self.next_input = self.next_input.half()
              # else:
              self.next_input = self.next_input.float()
              self.next_input = self.next_input.sub_(self.mean).div_(self.std)
  
      def next(self):
          torch.cuda.current_stream().wait_stream(self.stream)
          input = self.next_input
          target = self.next_target
          if input is not None:
              input.record_stream(torch.cuda.current_stream())
          if target is not None:
              target.record_stream(torch.cuda.current_stream())
          self.preload()
          return input, target
  
  
  def reduce_tensor(tensor, world_size):
      rt = tensor.clone()
      dist.all_reduce(rt, op=dist.reduce_op.SUM)
      rt /= world_size
      return rt
  
  
  
  def adjust_learning_rate(optimizer, epoch, step, len_epoch, args):
      """LR schedule that should yield 76% converged accuracy with batch size 256"""
      factor = epoch // 30
  
      if epoch >= 80:
          factor = factor + 1
  
      lr = args.learning_rate*(0.1**factor)
  
      """Warmup"""
      if epoch < 5:
          lr = lr*float(1 + step + epoch*len_epoch)/(5.*len_epoch)
  
      # if(args.apex_local_rank == 0):
      #     print("epoch = {}, step = {}, lr = {}".format(epoch, step, lr))
  
      for param_group in optimizer.param_groups:
          param_group['lr'] = lr
  
          
  # def adjust_learning_rate(optimizer, epoch, args):
  #     # Smaller slope for the last 5 epochs because lr * 1/250 is relatively large
  #     if args.epochs -  epoch > 5:
  #         lr = args.learning_rate * (args.epochs - 5 - epoch) / (args.epochs - 5)
  #     else:
  #         lr = args.learning_rate * (args.epochs - epoch) / ((args.epochs - 5) * 5)
  #     for param_group in optimizer.param_groups:
  #         param_group['lr'] = lr
  #     return lr    
  
  
  def train(train_loader, model, criterion, optimizer, epoch, args):
      batch_time = AverageMeter()
      losses = AverageMeter()
      top1 = AverageMeter()
      top5 = AverageMeter()
  
      # switch to train mode
      model.train()
      end = time.time()
      prefetcher = data_prefetcher(train_loader)
      input, target = prefetcher.next()
      i = 0
      while input is not None:
          i += 1
          if args.apex_profiling >= 0 and i == args.apex_profiling:
              print("Profiling begun at iteration {}".format(i))
              torch.cuda.cudart().cudaProfilerStart()
  
          if args.apex_profiling >= 0: torch.cuda.nvtx.range_push("Body of iteration {}".format(i))
  
          adjust_learning_rate(optimizer, epoch, i, len(train_loader), args)
  
          # compute output
          if args.apex_profiling >= 0: torch.cuda.nvtx.range_push("forward")
          logits, logtis_aux = model(input)
          if args.apex_profiling >= 0: torch.cuda.nvtx.range_pop()
          loss = criterion(logits, target)
          if args.auxiliary:
              loss_aux = criterion(logtis_aux, target)
              loss += args.auxiliary_weight * loss_aux
  
          # compute gradient and do SGD step
          optimizer.zero_grad()
  
          if args.apex_profiling >= 0: torch.cuda.nvtx.range_push("backward")
          with amp.scale_loss(loss, optimizer) as scaled_loss:
              scaled_loss.backward()
          if args.apex_profiling >= 0: torch.cuda.nvtx.range_pop()
  
          # for param in model.parameters():
          #     print(param.data.double().sum().item(), param.grad.data.double().sum().item())
  
          if args.apex_profiling >= 0: torch.cuda.nvtx.range_push("optimizer.step()")
          optimizer.step()
          if args.apex_profiling >= 0: torch.cuda.nvtx.range_pop()
  
          if i%args.report_freq == 0:
              # Every report_freq iterations, check the loss, accuracy, and speed.
              # For best performance, it doesn't make sense to print these metrics every
              # iteration, since they incur an allreduce and some host<->device syncs.
  
              # Measure accuracy
              prec1, prec5 = accuracy(logits.data, target, topk=(1, 5))
  
              # Average loss and accuracy across processes for logging
              if args.distributed:
                  reduced_loss = reduce_tensor(loss.data, args.world_size)
                  prec1 = reduce_tensor(prec1, args.world_size)
                  prec5 = reduce_tensor(prec5, args.world_size)
              else:
                  reduced_loss = loss.data
  
              # to_python_float incurs a host<->device sync
              losses.update(to_python_float(reduced_loss), input.size(0))
              top1.update(to_python_float(prec1), input.size(0))
              top5.update(to_python_float(prec5), input.size(0))
  
              torch.cuda.synchronize()
              batch_time.update((time.time() - end)/args.report_freq)
              end = time.time()
  
              if args.apex_local_rank == 0:
                  logging.info('Epoch: [{0}][{1}/{2}]\t'
                        'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                        'Speed {3:.3f} ({4:.3f})\t'
                        'Loss {loss.val:.10f} ({loss.avg:.4f})\t'
                        'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
                        'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
                         epoch, i, len(train_loader),
                         args.world_size*args.batch_size/batch_time.val,
                         args.world_size*args.batch_size/batch_time.avg,
                         batch_time=batch_time,
                         loss=losses, top1=top1, top5=top5))
          if args.apex_profiling >= 0: torch.cuda.nvtx.range_push("prefetcher.next()")
          input, target = prefetcher.next()
          if args.apex_profiling >= 0: torch.cuda.nvtx.range_pop()
  
          # Pop range "Body of iteration {}".format(i)
          if args.apex_profiling >= 0: torch.cuda.nvtx.range_pop()
  
          if args.apex_profiling >= 0 and i == args.apex_profiling + 10:
              print("Profiling ended at iteration {}".format(i))
              torch.cuda.cudart().cudaProfilerStop()
              quit()
      return top1.avg, losses.avg
  
  
  def validate(val_loader, model, criterion, args):
      batch_time = AverageMeter()
      losses = AverageMeter()
      top1 = AverageMeter()
      top5 = AverageMeter()
  
      # switch to evaluate mode
      model.eval()
  
      end = time.time()
  
      prefetcher = data_prefetcher(val_loader)
      input, target = prefetcher.next()
      i = 0
      while input is not None:
          i += 1
  
          # compute output
          with torch.no_grad():
              output, _ = model(input)
              loss = criterion(output, target)
  
          # measure accuracy and record loss
          prec1, prec5 = accuracy(output.data, target, topk=(1, 5))
  
          if args.distributed:
              reduced_loss = reduce_tensor(loss.data, args.world_size)
              prec1 = reduce_tensor(prec1, args.world_size)
              prec5 = reduce_tensor(prec5, args.world_size)
          else:
              reduced_loss = loss.data
  
          losses.update(to_python_float(reduced_loss), input.size(0))
          top1.update(to_python_float(prec1), input.size(0))
          top5.update(to_python_float(prec5), input.size(0))
  
          # measure elapsed time
          batch_time.update(time.time() - end)
          end = time.time()
  
          # TODO:  Change timings to mirror train().
          if args.apex_local_rank == 0 and i % args.report_freq == 0:
              logging.info('Test: [{0}/{1}]\t'
                    'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                    'Speed {2:.3f} ({3:.3f})\t'
                    'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                    'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
                    'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
                     i, len(val_loader),
                     args.world_size * args.batch_size / batch_time.val,
                     args.world_size * args.batch_size / batch_time.avg,
                     batch_time=batch_time, loss=losses,
                     top1=top1, top5=top5))
  
          input, target = prefetcher.next()
  
      logging.info(' * Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f}'
            .format(top1=top1, top5=top5))
  
      return top1.avg, top5.avg, losses.avg
  
  
  def dali_apex_train(train_loader, model, criterion, optimizer, epoch, args):
      batch_time = AverageMeter()
      losses = AverageMeter()
      top1 = AverageMeter()
      top5 = AverageMeter()
  
      # switch to train mode
      model.train()
      end = time.time()
  
      for i, data in enumerate(train_loader):
          input = data[0]["data"]
          target = data[0]["label"].squeeze().cuda().long()
          train_loader_len = int(math.ceil(train_loader._size / args.batch_size))
  
          if args.dali_profiling >= 0 and i == args.dali_profiling:
              print("Profiling begun at iteration {}".format(i))
              torch.cuda.cudart().cudaProfilerStart()
  
          if args.dali_profiling >= 0: torch.cuda.nvtx.range_push("Body of iteration {}".format(i))
  
          # adjust_learning_rate(optimizer, epoch, i, train_loader_len, args)
          if args.debug:
              if i > 10:
                  logging.info('Break in debug mode after 10 batchs...')
                  break
  
          # compute output
          if args.dali_profiling >= 0: torch.cuda.nvtx.range_push("forward")
          logits, logtis_aux = model(input)
          if args.dali_profiling >= 0: torch.cuda.nvtx.range_pop()
          loss = criterion(logits, target)
          if args.auxiliary:
              loss_aux = criterion(logtis_aux, target)
              loss += args.auxiliary_weight * loss_aux
  
          # compute gradient and do SGD step
          optimizer.zero_grad()
  
          if args.dali_profiling >= 0: torch.cuda.nvtx.range_push("backward")
          if args.apex_opt_level is not None:
              with amp.scale_loss(loss, optimizer) as scaled_loss:
                  scaled_loss.backward()
          else:
               loss.backward()
          if args.dali_profiling >= 0: torch.cuda.nvtx.range_pop()
  
          if args.dali_profiling >= 0: torch.cuda.nvtx.range_push("optimizer.step()")
          optimizer.step()
          if args.dali_profiling >= 0: torch.cuda.nvtx.range_pop()
  
          if i%args.report_freq == 0:
              # Every print_freq iterations, check the loss, accuracy, and speed.
              # For best performance, it doesn't make sense to print these metrics every
              # iteration, since they incur an allreduce and some host<->device syncs.
  
              # Measure accuracy
              prec1, prec5 = accuracy(logits.data, target, topk=(1, 5))
  
              # Average loss and accuracy across processes for logging
              if args.distributed:
                  reduced_loss = reduce_tensor(loss.data, args.world_size)
                  prec1 = reduce_tensor(prec1, args.world_size)
                  prec5 = reduce_tensor(prec5, args.world_size)
              else:
                  reduced_loss = loss.data
  
              # to_python_float incurs a host<->device sync
              losses.update(to_python_float(reduced_loss), input.size(0))
              top1.update(to_python_float(prec1), input.size(0))
              top5.update(to_python_float(prec5), input.size(0))
  
              torch.cuda.synchronize()
              batch_time.update((time.time() - end)/args.report_freq)
              end = time.time()
  
              if args.apex_local_rank == 0:
                  logging.info('Epoch: [{0}][{1}/{2}]\t'
                        'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                        'Speed {3:.3f} ({4:.3f})\t'
                        'Loss {loss.val:.10f} ({loss.avg:.4f})\t'
                        'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
                        'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
                         epoch, i, train_loader_len,
                         args.world_size*args.batch_size/batch_time.val,
                         args.world_size*args.batch_size/batch_time.avg,
                         batch_time=batch_time,
                         loss=losses, top1=top1, top5=top5))
  
          # Pop range "Body of iteration {}".format(i)
          if args.dali_profiling >= 0: torch.cuda.nvtx.range_pop()
  
          if args.dali_profiling >= 0 and i == args.dali_profiling + 2:
              print("Profiling ended at iteration {}".format(i))
              torch.cuda.cudart().cudaProfilerStop()
              quit()
  
      return top1.avg, losses.avg
  
  
  
  def dali_validate(val_loader, model, criterion, args):
      batch_time = AverageMeter()
      losses = AverageMeter()
      top1 = AverageMeter()
      top5 = AverageMeter()
  
      # switch to evaluate mode
      model.eval()
  
      end = time.time()
  
      for i, data in enumerate(val_loader):
          input = data[0]["data"]
          target = data[0]["label"].squeeze().cuda().long()
          val_loader_len = int(val_loader._size / args.batch_size)
  
          if args.debug:
              if i > 10:
                  break
  
          # compute output
          with torch.no_grad():
              output, _ = model(input)
              loss = criterion(output, target)
  
          # measure accuracy and record loss
          prec1, prec5 = accuracy(output.data, target, topk=(1, 5))
  
          if args.distributed:
              reduced_loss = reduce_tensor(loss.data, args.world_size)
              prec1 = reduce_tensor(prec1, args.world_size)
              prec5 = reduce_tensor(prec5, args.world_size)
          else:
              reduced_loss = loss.data
  
          losses.update(to_python_float(reduced_loss), input.size(0))
          top1.update(to_python_float(prec1), input.size(0))
          top5.update(to_python_float(prec5), input.size(0))
  
          # measure elapsed time
          batch_time.update(time.time() - end)
          end = time.time()
  
          # TODO:  Change timings to mirror train().
          if args.apex_local_rank == 0 and i % args.report_freq == 0:
              logging.info('Test: [{0}/{1}]\t'
                    'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                    'Speed {2:.3f} ({3:.3f})\t'
                    'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                    'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
                    'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
                     i, val_loader_len,
                     args.world_size * args.batch_size / batch_time.val,
                     args.world_size * args.batch_size / batch_time.avg,
                     batch_time=batch_time, loss=losses,
                     top1=top1, top5=top5))
  
      logging.info(' * Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f}'
          .format(top1=top1, top5=top5))
  
      return top1.avg, top5.avg, losses.avg
  
apis: torch.cuda.cudart torch.no_grad torch.tensor numpy.asarray torch.cuda.synchronize torch.from_numpy torch.distributed.all_reduce numpy.expand_dims torch.cuda.nvtx.range_push torch.cuda.current_stream torch.cuda.stream torch.cuda.Stream torch.cuda.nvtx.range_pop