text: import os
  import sys
  import time
  import logging
  import torch
  import utils as dutils
  import argparse
  import numpy as np
  import torch.utils
  from genotypes import Genotype_nested
  import genotypes
  import torch.nn as nn
  import torchvision.datasets as dset
  import torch.backends.cudnn as cudnn
  from model import NetworkCIFAR, NetworkImageNet
  from thop import profile
  from torch.utils.tensorboard import SummaryWriter
  
  parser = argparse.ArgumentParser("cifar")
  parser.add_argument('--data', type=str, default=os.path.join(os.path.expanduser('~'),'work/dataset/cifar/'), help='location of the data corpus')
  parser.add_argument('--dataset', type=str, default='cifar10', help='cifar10 or cifar100 or imagenet or coco_captions or coco_detection for training')
  parser.add_argument('--batch_size', type=int, default=128, help='batch size')
  parser.add_argument('--learning_rate', type=float, default=0.025, help='init learning rate')
  parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
                      help='manual epoch number (useful on restarts)')
  parser.add_argument('--momentum', type=float, default=0.9, help='momentum')
  parser.add_argument('--weight_decay', type=float, default=3e-4, help='weight decay')
  parser.add_argument('--report_freq', type=int, default=50, help='report frequency')
  parser.add_argument('--gpu', type=int, default=0, help='gpu device id')
  parser.add_argument('--multi-gpus', action='store_true', default=False, help='use multi gpus')
  parser.add_argument('--parse_method', type=str, default='threshold_sparse', help='experiment name')
  parser.add_argument('--epochs', type=int, default=600, help='num of training epochs')
  parser.add_argument('--init_channels', type=int, default=64, help='num of init channels')
  parser.add_argument('--layers', type=int, default=8, help='total number of layers')
  parser.add_argument('--model_path', type=str, default='saved_models', help='path to save the model')
  parser.add_argument('--auxiliary', action='store_true', default=False, help='use auxiliary tower')
  parser.add_argument('--auxiliary_weight', type=float, default=0.4, help='weight for auxiliary loss')
  parser.add_argument('--cutout', action='store_true', default=False, help='use cutout')
  parser.add_argument('--cutout_length', type=int, default=16, help='cutout length')
  parser.add_argument('--label_smooth', type=float, default=0.1, help='label smoothing')
  parser.add_argument('--auto_aug', action='store_true', default=False, help='use auto augmentation')
  parser.add_argument('--drop_path_prob', type=float, default=0.2, help='drop path probability')
  parser.add_argument('--save', type=str, default='EXP', help='experiment name')
  parser.add_argument('--seed', type=int, default=0, help='random seed')
  parser.add_argument('--arch', type=str, default='nestedDARTS_cell_threshold_sparse_0.85', help='which architecture to use (whole string or txt file location)')
  parser.add_argument('--grad_clip', type=float, default=5, help='gradient clipping')
  parser.add_argument('--resume', default='', type=str, metavar='PATH', help='path to latest checkpoint (default: none)')
  parser.add_argument('--amp', action='store_true', default=False, help="Train using Automatic Mixed Precision.")
  parser.add_argument('--no_log', action='store_true', help="Do not record training log.")
  
  
  class CrossEntropyLabelSmooth(nn.Module):
      "Smooth CrossEntropy loss for training on ImageNet."
  
      def __init__(self, num_classes, epsilon):
          super(CrossEntropyLabelSmooth, self).__init__()
          self.num_classes = num_classes
          self.epsilon = epsilon
          self.logsoftmax = nn.LogSoftmax(dim=1)
  
      def forward(self, inputs, targets):
          log_probs = self.logsoftmax(inputs)
          targets = torch.zeros_like(log_probs).scatter_(1, targets.unsqueeze(1), 1)
          targets = (1 - self.epsilon) * targets + self.epsilon / self.num_classes
          loss = (-targets * log_probs).mean(0).sum()
          return loss
  
  class TrainArgs:
  
      def __init__(self, arch, epochs, dataset, init_channels, layers, gpu, no_log=True) -> None:
          self.arch = arch
          self.dataset = dataset
          self.epochs = epochs    
          self.init_channels = init_channels
          self.layers = layers
          self.no_log = no_log
          self.data = os.path.join(os.path.expanduser('~'),'work/dataset/cifar/')
          self.batch_size = 256
          self.learning_rate = 0.025
          self.start_epoch = 0
          self.momentum = 0.9
          self.weight_decay = 3e-4
          self.report_freq = 100
          self.gpu = gpu
          self.multi_gpus = False
          self.parse_method = "threshold_sparse"
          self.model_path = "saved_models"
          self.auxiliary = False
          self.auxiliary_weight = 0.4
          self.cutout = False
          self.cutout_length = 16
          self.label_smooth = 0.1
          self.auto_aug = False
          self.drop_path_prob = 0.2
          self.save = 'EXP'
          self.seed = 0
          self.grad_clip = 5
          self.resume = ''
          self.amp = True
  
          
  class TrainNetwork(object):
      """The main train network"""
  
      def __init__(self, args):
          super(TrainNetwork, self).__init__()
          self.args = args
          self.dur_time = 0
          self._init_log()
          self._init_device()
          self._init_data_queue()
          self._init_model()
  
      def _init_log(self):
          if not self.args.no_log:
              self.args.save = './logs/eval/' + self.args.arch + '/' + 'cifar10' + '/eval-{}-{}'.format(self.args.save, time.strftime('%Y%m%d-%H%M'))
              dutils.create_exp_dir(self.args.save, scripts_to_save=None)
  
          log_format = '%(asctime)s %(message)s'
          logging.basicConfig(stream=sys.stdout, level=logging.INFO,
                              format=log_format, datefmt='%m/%d %I:%M:%S %p')
          self.logger = logging.getLogger('Architecture Training')
          if not self.args.no_log:
              fh = logging.FileHandler(os.path.join(self.args.save, 'log.txt'))
              fh.setFormatter(logging.Formatter(log_format))
         
              self.logger.addHandler(fh)
              self.writer = SummaryWriter(f'./runs/eval/eval-{args.dataset}-{args.layers}layers-{time.strftime("%Y%m%d-%H%M")}')
  
      def _init_device(self):
          if not torch.cuda.is_available():
              self.logger.info('no gpu device available')
              sys.exit(1)
          np.random.seed(self.args.seed)
          self.device_id = self.args.gpu
          self.device = torch.device('cuda:{}'.format(0 if self.args.multi_gpus else self.device_id))
          cudnn.benchmark = True
          torch.manual_seed(self.args.seed)
          cudnn.enabled = True
          torch.cuda.manual_seed(self.args.seed)
          logging.info('gpu device = %d' % self.args.gpu)
          logging.info("args = %s", self.args)
  
      def _init_data_queue(self):
          if self.args.dataset == 'cifar10':
              train_transform, valid_transform = dutils._data_transforms_cifar(self.args)
              train_data = dset.CIFAR10(root=self.args.data, train=True, download=True, transform=train_transform)
              valid_data = dset.CIFAR10(root=self.args.data, train=False, download=True, transform=valid_transform)
              self.num_classes = 10
          elif self.args.dataset == 'cifar100':
              train_transform, valid_transform = dutils._data_transforms_cifar(self.args)
              train_data = dset.CIFAR100(root=self.args.data, train=True, download=True, transform=train_transform)
              valid_data = dset.CIFAR100(root=self.args.data, train=False, download=True, transform=valid_transform)
              self.num_classes = 100
          elif self.args.dataset == 'imagenet':
              train_transform, valid_transform = dutils._data_transforms_imagenet(self.args)
              train_data = dset.ImageNet(root=self.args.data, split='train', transform=train_transform)
              valid_data = dset.ImageNet(root=self.args.data, split='val', transform=valid_transform)
              self.num_classes = 1000
          self.train_queue = torch.utils.data.DataLoader(
              train_data, batch_size=self.args.batch_size, shuffle=True, pin_memory=True, num_workers=4)
  
          self.valid_queue = torch.utils.data.DataLoader(
              valid_data, batch_size=self.args.batch_size, shuffle=False, pin_memory=True, num_workers=4)
  
      def _init_model(self):
          if "baseline" in self.args.arch:
              self.genotype = eval(f"genotypes.{self.args.arch.split('_')[1]}")
          elif isinstance(self.args.arch, Genotype_nested):
              self.genotype = self.args.arch
          else:
              with open(f'genotypes/{self.args.arch}.txt', 'r') as g:
                  self.genotype = eval(g.read())
          if self.args.dataset == "imagenet":
              model = NetworkImageNet(self.args.init_channels, self.num_classes, self.args.layers, self.args.auxiliary, self.genotype, self.args.parse_method)
              inputs_r = (torch.randn(1, 3, 224, 224),)
          else:
              model = NetworkCIFAR(self.args.init_channels, self.num_classes, self.args.layers, self.args.auxiliary, self.genotype, self.args.parse_method)
              inputs_r = (torch.randn(1, 3, 224, 224),)
          flops, params = profile(model, inputs=inputs_r, verbose=False)
          self.logger.info('flops = %fM', flops / 1e6)
          self.logger.info('param size = %fM', params / 1e6)
          
          # Try move model to multi gpus
          if torch.cuda.device_count() > 1 and self.args.multi_gpus:
              self.logger.info('use: %d gpus', torch.cuda.device_count())
              model = nn.DataParallel(model)
          else:
              self.logger.info('gpu device = %d' % self.device_id)
              torch.cuda.set_device(self.device_id)
          self.model = model.to(self.device)
  
          criterion = nn.CrossEntropyLoss()
          criterion_smooth = CrossEntropyLabelSmooth(self.num_classes, self.args.label_smooth)
          self.criterion = criterion.to(self.device)
          self.criterion_train = criterion_smooth.to(self.device) if self.args.dataset == "imagenet" else self.criterion
          self.optimizer = torch.optim.SGD(
              model.parameters(),
              self.args.learning_rate,
              momentum=self.args.momentum,
              weight_decay=self.args.weight_decay
          )
  
          self.scaler = torch.cuda.amp.GradScaler(enabled=self.args.amp)
  
          self.best_acc_top1 = 0
          # optionally resume from a checkpoint
          if self.args.resume:
              if os.path.isfile(self.args.resume):
                  print("=> loading checkpoint {}".format(self.args.resume))
                  checkpoint = torch.load(self.args.resume, map_location=self.device)
                  self.dur_time = checkpoint['dur_time']
                  self.args.start_epoch = checkpoint['epoch']
                  self.best_acc_top1 = checkpoint['best_acc_top1']
                  self.args.drop_path_prob = checkpoint['drop_path_prob']
                  self.model.load_state_dict(checkpoint['state_dict'])
                  self.optimizer.load_state_dict(checkpoint['optimizer'])
                  self.scaler.load_state_dict(checkpoint['scaler'])
                  print("=> loaded checkpoint '{}' (epoch {})".format(self.args.resume, checkpoint['epoch']))
              else:
                  print("=> no checkpoint found at '{}'".format(self.args.resume))
  
          last_epoch = -1 if self.args.start_epoch == 0 else self.args.start_epoch
          self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.optimizer, float(self.args.epochs), eta_min=0, last_epoch=last_epoch)
          # reload the scheduler if possible
          if self.args.resume and os.path.isfile(self.args.resume):
              checkpoint = torch.load(self.args.resume)
              self.scheduler.load_state_dict(checkpoint['scheduler'])
  
      def run(self):
          self.logger.info('args = %s', self.args)
          run_start = time.time()
          for epoch in range(self.args.start_epoch, self.args.epochs):
              lr = self.scheduler.get_lr()[0]
              self.logger.info('epoch % d / %d  lr %e', epoch, self.args.epochs, lr)
  
              self.model.drop_path_prob = self.args.drop_path_prob * epoch / self.args.epochs
  
              train_acc, train_obj = self.train()
              self.logger.info('train loss %e, train acc %f', train_obj, train_acc)
  
              valid_acc_top1, valid_acc_top5, valid_obj = self.infer()
              self.logger.info('valid loss %e, top1 valid acc %f top5 valid acc %f',
                          valid_obj, valid_acc_top1, valid_acc_top5)
              self.logger.info('best valid acc %f', self.best_acc_top1)
  
              self.scheduler.step()
  
              is_best = False
              if valid_acc_top1 > self.best_acc_top1:
                  self.best_acc_top1 = valid_acc_top1
                  is_best = True
              
              if not self.args.no_log:
                  self.writer.add_scalar('lr', lr, global_step=epoch)
                  self.writer.add_scalar('valid acc top1', valid_acc_top1, global_step=epoch)
                  self.writer.add_scalar('valid acc top5', valid_acc_top5, global_step=epoch)
                  self.writer.add_scalar('train acc top1', train_acc, global_step=epoch)
  
                  dutils.save_checkpoint({
                      'epoch': epoch+1,
                      'dur_time': self.dur_time + time.time() - run_start,
                      'state_dict': self.model.state_dict(),
                      'drop_path_prob': self.args.drop_path_prob,
                      'best_acc_top1': self.best_acc_top1,
                      'optimizer': self.optimizer.state_dict(),
                      'scaler': self.scaler.state_dict(),
                      'scheduler': self.scheduler.state_dict()
                  }, is_best, self.args.save)
          
          duration_time = dutils.calc_time(self.dur_time + time.time() - run_start)
          self.logger.info('train epoches %d, best_acc_top1 %f, dur_time %s',
                           self.args.epochs, self.best_acc_top1, duration_time)
          
          if not self.args.no_log:
              self.writer.close()
  
          return {'type': "Eval", 'val_acc': self.best_acc_top1, "model": self.args.genotype, "dataset": self.args.dataset, "duration_time": duration_time}
  
      def train(self):
          objs = dutils.AvgrageMeter()
          top1 = dutils.AvgrageMeter()
          top5 = dutils.AvgrageMeter()
  
          self.model.train()
  
          for step, (input, target) in enumerate(self.train_queue):
              input = input.cuda(non_blocking=True)
              target = target.cuda(non_blocking=True)
  
              self.optimizer.zero_grad(set_to_none=True)
  
              with torch.cuda.amp.autocast(enabled=self.args.amp):
                  logits, logits_aux = self.model(input)
                  loss = self.criterion_train(logits, target)
                  if self.args.auxiliary:
                      loss_aux = self.criterion_train(logits_aux, target)
                      loss += self.args.auxiliary_weight*loss_aux
  
              self.scaler.scale(loss).backward()
              self.scaler.unscale_(self.optimizer)
              nn.utils.clip_grad_norm_(self.model.parameters(), self.args.grad_clip)
              self.scaler.step(self.optimizer)
              self.scaler.update()
  
              prec1, prec5 = dutils.accuracy(logits, target, topk=(1, 5))
              n = input.size(0)
              objs.update(loss.item(), n)
              top1.update(prec1.item(), n)
              top5.update(prec5.item(), n)
  
              if step % self.args.report_freq == 0:
                  self.logger.info('train %03d %e %f %f', step, objs.avg, top1.avg, top5.avg)
  
          return top1.avg, objs.avg
  
      def infer(self):
          objs = dutils.AvgrageMeter()
          top1 = dutils.AvgrageMeter()
          top5 = dutils.AvgrageMeter()
          self.model.eval()
          with torch.no_grad():
              for step, (input, target) in enumerate(self.valid_queue):
                  input = input.cuda(non_blocking=True)
                  target = target.cuda(non_blocking=True)
  
                  with torch.cuda.amp.autocast(enabled=self.args.amp):
                      logits, _ = self.model(input)
                      loss = self.criterion(logits, target)
  
                  prec1, prec5 = dutils.accuracy(logits, target, topk=(1, 5))
                  n = input.size(0)
                  objs.update(loss.item(), n)
                  top1.update(prec1.item(), n)
                  top5.update(prec5.item(), n)
  
                  if step % self.args.report_freq == 0:
                      self.logger.info('valid %03d %e %f %f', step, objs.avg, top1.avg, top5.avg)
              return top1.avg, top5.avg, objs.avg
  
  
  if __name__ == '__main__':
      args = parser.parse_args()
      train_network = TrainNetwork(args)
      train_network.run()