text: import os
  import time
  
  import numpy as np
  import torch.nn as nn
  from tqdm import tqdm
  from collections import OrderedDict
  
  from dataloaders.datasets.hsi_dataset import get_training_dataloaders
  from mypath import Path
  from dataloaders import make_data_loader
  from modeling.sync_batchnorm.replicate import patch_replication_callback
  from modeling.deeplab import *
  from utils.loss import SegmentationLosses
  from utils.lr_scheduler import LR_Scheduler
  from utils.saver import Saver
  from utils.summaries import TensorboardSummary
  from utils.metrics import Evaluator
  from auto_deeplab import AutoDeeplab
  from config_utils.search_args import obtain_search_args
  from utils.copy_state_dict import copy_state_dict
  import apex
  try:
      from apex import amp
      APEX_AVAILABLE = True
  except ModuleNotFoundError:
      APEX_AVAILABLE = False
  
  
  print('working with pytorch version {}'.format(torch.__version__))
  print('with cuda version {}'.format(torch.version.cuda))
  print('cudnn enabled: {}'.format(torch.backends.cudnn.enabled))
  print('cudnn version: {}'.format(torch.backends.cudnn.version()))
  
  torch.backends.cudnn.benchmark = True
  
  class Trainer(object):
      def __init__(self, args):
          self.args = args
  
          # Define Saver
          self.saver = Saver(args)
          self.saver.save_experiment_config()
          # Define Tensorboard Summary
          self.summary = TensorboardSummary(self.saver.experiment_dir)
          self.writer = self.summary.create_summary()
          self.use_amp = True if (APEX_AVAILABLE and args.use_amp) else False
          self.opt_level = args.opt_level
  
          kwargs = {'num_workers': args.workers, 'pin_memory': True, 'drop_last':True}
          self.train_loaderA, self.train_loaderB, self.val_loader = get_training_dataloaders(args.batch_size, args.workers, '/projects/datasets/UOW-HSI-v2', fold=args.fold)
          self.nclass=5
  
          if args.use_balanced_weights:
              classes_weights_path = os.path.join(Path.db_root_dir(args.dataset), args.dataset+'_classes_weights.npy')
              if os.path.isfile(classes_weights_path):
                  weight = np.load(classes_weights_path)
              else:
                  raise NotImplementedError
                  #if so, which trainloader to use?
                  # weight = calculate_weigths_labels(args.dataset, self.train_loader, self.nclass)
              weight = torch.from_numpy(weight.astype(np.float32))
          else:
              weight = None
          self.criterion = SegmentationLosses(weight=weight, cuda=args.cuda).build_loss(mode=args.loss_type)
  
          # Define network
          model = AutoDeeplab (self.nclass, 10, self.criterion, self.args.filter_multiplier,
                               self.args.block_multiplier, self.args.step)
          optimizer = torch.optim.SGD(
                  model.weight_parameters(),
                  args.lr,
                  momentum=args.momentum,
                  weight_decay=args.weight_decay
              )
  
          self.model, self.optimizer = model, optimizer
  
          self.architect_optimizer = torch.optim.Adam(self.model.arch_parameters(),
                                                      lr=args.arch_lr, betas=(0.9, 0.999),
                                                      weight_decay=args.arch_weight_decay)
  
          # Define Evaluator
          self.evaluator = Evaluator(self.nclass)
          # Define lr scheduler
          self.scheduler = LR_Scheduler(args.lr_scheduler, args.lr,
                                        args.epochs, len(self.train_loaderA), min_lr=args.min_lr)
          # TODO: Figure out if len(self.train_loader) should be devided by two ? in other module as well
          # Using cuda
          if args.cuda:
              self.model = self.model.cuda()
  
  
          # mixed precision
          if self.use_amp and args.cuda:
              keep_batchnorm_fp32 = True if (self.opt_level == 'O2' or self.opt_level == 'O3') else None
  
              # fix for current pytorch version with opt_level 'O1'
              if self.opt_level == 'O1' and torch.__version__ < '1.3':
                  for module in self.model.modules():
                      if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
                          # Hack to fix BN fprop without affine transformation
                          if module.weight is None:
                              module.weight = torch.nn.Parameter(
                                  torch.ones(module.running_var.shape, dtype=module.running_var.dtype,
                                             device=module.running_var.device), requires_grad=False)
                          if module.bias is None:
                              module.bias = torch.nn.Parameter(
                                  torch.zeros(module.running_var.shape, dtype=module.running_var.dtype,
                                              device=module.running_var.device), requires_grad=False)
  
              # print(keep_batchnorm_fp32)
              self.model, [self.optimizer, self.architect_optimizer] = amp.initialize(
                  self.model, [self.optimizer, self.architect_optimizer], opt_level=self.opt_level,
                  keep_batchnorm_fp32=keep_batchnorm_fp32, loss_scale="dynamic")
  
              print('cuda finished')
  
  
          # Using data parallel
          if args.cuda and len(self.args.gpu_ids) >1:
              if self.opt_level == 'O2' or self.opt_level == 'O3':
                  print('currently cannot run with nn.DataParallel and optimization level', self.opt_level)
              self.model = torch.nn.DataParallel(self.model, device_ids=self.args.gpu_ids)
              patch_replication_callback(self.model)
              print('training on multiple-GPUs')
  
          #checkpoint = torch.load(args.resume)
          #print('about to load state_dict')
          #self.model.load_state_dict(checkpoint['state_dict'])
          #print('model loaded')
          #sys.exit()
  
          # Resuming checkpoint
          self.best_pred = 0.0
          if args.resume is not None:
              if not os.path.isfile(args.resume):
                  raise RuntimeError("=> no checkpoint found at '{}'" .format(args.resume))
              checkpoint = torch.load(args.resume)
              args.start_epoch = checkpoint['epoch']
  
              # if the weights are wrapped in module object we have to clean it
              if args.clean_module:
                  self.model.load_state_dict(checkpoint['state_dict'])
                  state_dict = checkpoint['state_dict']
                  new_state_dict = OrderedDict()
                  for k, v in state_dict.items():
                      name = k[7:]  # remove 'module.' of dataparallel
                      new_state_dict[name] = v
                  # self.model.load_state_dict(new_state_dict)
                  copy_state_dict(self.model.state_dict(), new_state_dict)
  
              else:
                  if torch.cuda.device_count() > 1 or args.load_parallel:
                      # self.model.module.load_state_dict(checkpoint['state_dict'])
                      copy_state_dict(self.model.module.state_dict(), checkpoint['state_dict'])
                  else:
                      # self.model.load_state_dict(checkpoint['state_dict'])
                      copy_state_dict(self.model.state_dict(), checkpoint['state_dict'])
  
  
              if not args.ft:
                  # self.optimizer.load_state_dict(checkpoint['optimizer'])
                  copy_state_dict(self.optimizer.state_dict(), checkpoint['optimizer'])
              self.best_pred = checkpoint['best_pred']
              print("=> loaded checkpoint '{}' (epoch {})"
                    .format(args.resume, checkpoint['epoch']))
  
          # Clear start epoch if fine-tuning
          if args.ft:
              args.start_epoch = 0
  
      def training(self, epoch):
          train_loss = 0.0
          self.model.train()
          tbar = tqdm(self.train_loaderA)
          num_img_tr = len(self.train_loaderA)
          for i, sample in enumerate(tbar):
              image, target = sample['image'], sample['label']
              if len(image) <= 1:  # skip batches that are not more than 1
                  continue
              if self.args.cuda:
                  image, target = image.cuda(), target.cuda()
              self.scheduler(self.optimizer, i, epoch, self.best_pred)
              self.optimizer.zero_grad()
              output = self.model(image)
              loss = self.criterion(output, target)
              if self.use_amp:
                  with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                      scaled_loss.backward()
              else:
                  loss.backward()
              self.optimizer.step()
  
              if epoch >= self.args.alpha_epoch:
                  search = next(iter(self.train_loaderB))
                  image_search, target_search = search['image'], search['label']
                  if self.args.cuda:
                      image_search, target_search = image_search.cuda (), target_search.cuda ()
  
                  self.architect_optimizer.zero_grad()
                  output_search = self.model(image_search)
                  arch_loss = self.criterion(output_search, target_search)
                  if self.use_amp:
                      with amp.scale_loss(arch_loss, self.architect_optimizer) as arch_scaled_loss:
                          arch_scaled_loss.backward()
                  else:
                      arch_loss.backward()
                  self.architect_optimizer.step()
  
              train_loss += loss.item()
              tbar.set_description('Train loss: %.3f' % (train_loss / (i + 1)))
              #self.writer.add_scalar('train/total_loss_iter', loss.item(), i + num_img_tr * epoch)
  
              # Show 10 * 3 inference results each epoch
              # if i % (num_img_tr // 10) == 0:
              #     global_step = i + num_img_tr * epoch
                  # self.summary.visualize_image(self.writer, self.args.dataset, image, target, output, global_step)
  
              #torch.cuda.empty_cache()
          self.writer.add_scalar('train/total_loss_epoch', train_loss, epoch)
          print('[Epoch: %d, numImages: %5d]' % (epoch, i * self.args.batch_size + image.data.shape[0]))
          print('Loss: %.3f' % train_loss)
  
          if self.args.no_val:
              # save checkpoint every epoch
              is_best = False
              if torch.cuda.device_count() > 1:
                  state_dict = self.model.module.state_dict()
              else:
                  state_dict = self.model.state_dict()
              self.saver.save_checkpoint({
                  'epoch': epoch + 1,
                  'state_dict': state_dict,
                  'optimizer': self.optimizer.state_dict(),
                  'best_pred': self.best_pred,
              }, is_best)
  
  
      def validation(self, epoch):
          self.model.eval()
          self.evaluator.reset()
          tbar = tqdm(self.val_loader, desc='\r')
          test_loss = 0.0
  
          for i, sample in enumerate(tbar):
              image, target = sample['image'], sample['label']
              if self.args.cuda:
                  image, target = image.cuda(), target.cuda()
              with torch.no_grad():
                  output = self.model(image)
              loss = self.criterion(output, target)
              test_loss += loss.item()
              tbar.set_description('Test loss: %.3f' % (test_loss / (i + 1)))
              pred = output.data.cpu().numpy()
              target = target.cpu().numpy()
              pred = np.argmax(pred, axis=1)
              # Add batch sample into evaluator
              self.evaluator.add_batch(target, pred)
  
          # Fast test during the training
          Acc = self.evaluator.Pixel_Accuracy()
          Acc_class = self.evaluator.Pixel_Accuracy_Class()
          mIoU = self.evaluator.Mean_Intersection_over_Union()
          FWIoU = self.evaluator.Frequency_Weighted_Intersection_over_Union()
          self.writer.add_scalar('val/total_loss_epoch', test_loss, epoch)
          self.writer.add_scalar('val/mIoU', mIoU, epoch)
          self.writer.add_scalar('val/Acc', Acc, epoch)
          self.writer.add_scalar('val/Acc_class', Acc_class, epoch)
          self.writer.add_scalar('val/fwIoU', FWIoU, epoch)
          print('Validation:')
          print('[Epoch: %d, numImages: %5d]' % (epoch, i * self.args.batch_size + image.data.shape[0]))
          print("Acc:{}, Acc_class:{}, mIoU:{}, fwIoU: {}".format(Acc, Acc_class, mIoU, FWIoU))
          print('Loss: %.3f' % test_loss)
          new_pred = mIoU
          if new_pred > self.best_pred:
              is_best = True
              self.best_pred = new_pred
              if torch.cuda.device_count() > 1:
                  state_dict = self.model.module.state_dict()
              else:
                  state_dict = self.model.state_dict()
              self.saver.save_checkpoint({
                  'epoch': epoch + 1,
                  'state_dict': state_dict,
                  'optimizer': self.optimizer.state_dict(),
                  'best_pred': self.best_pred,
              }, is_best)
  
  def main():
      args = obtain_search_args()
      args.cuda = not args.no_cuda and torch.cuda.is_available()
      if args.cuda:
          try:
              args.gpu_ids = [int(s) for s in args.gpu_ids.split(',')]
          except ValueError:
              raise ValueError('Argument --gpu_ids must be a comma-separated list of integers only')
  
      if args.sync_bn is None:
          if args.cuda and len(args.gpu_ids) > 1:
              args.sync_bn = True
          else:
              args.sync_bn = False
  
      # default settings for epochs, batch_size and lr
      if args.epochs is None:
          epoches = {
              'coco': 30,
              'cityscapes': 40,
              'pascal': 50,
              'kd':10
          }
          args.epochs = epoches[args.dataset.lower()]
  
      if args.batch_size is None:
          args.batch_size = 4 * len(args.gpu_ids)
  
      if args.test_batch_size is None:
          args.test_batch_size = args.batch_size
  
      #args.lr = args.lr / (4 * len(args.gpu_ids)) * args.batch_size
  
  
      if args.checkname is None:
          args.checkname = 'deeplab-'+str(args.backbone)
      print(args)
      torch.manual_seed(args.seed)
      trainer = Trainer(args)
      print('Starting Epoch:', trainer.args.start_epoch)
      print('Total Epoches:', trainer.args.epochs)
      for epoch in range(trainer.args.start_epoch, trainer.args.epochs):
          trainer.training(epoch)
          if not trainer.args.no_val and epoch % args.eval_interval == (args.eval_interval - 1):
              trainer.validation(epoch)
  
      trainer.writer.close()
  
  if __name__ == "__main__":
      start_t = time.perf_counter()
      main()
      elapsed_time = time.perf_counter() - start_t
      print(elapsed_time/60/60)
  