text: import os
  import sys
  import glob
  import time
  import copy
  import random
  import numpy as np
  import torch
  import utils
  import logging
  import argparse
  import torch.nn as nn
  import torch.utils
  import torch.nn.functional as F
  import torchvision.datasets as dset
  import torchvision.transforms as transforms
  import torch.backends.cudnn as cudnn
  from model import NASNetworkCIFAR, NASNetworkImageNet
  from model_search import NASWSNetworkCIFAR, NASWSNetworkImageNet
  from controller import NAO
  
  parser = argparse.ArgumentParser(description='NAO CIFAR-10')
  
  # Basic model parameters.
  parser.add_argument('--mode', type=str, default='train', choices=['train', 'test'])
  parser.add_argument('--data', type=str, default='./data')
  parser.add_argument('--dataset', type=str, default='cifar10', choices=['cifar10, cifar100, imagenet'])
  parser.add_argument('--zip_file', action='store_true', default=False)
  parser.add_argument('--lazy_load', action='store_true', default=False)
  parser.add_argument('--output_dir', type=str, default='models')
  parser.add_argument('--seed', type=int, default=0)
  parser.add_argument('--child_batch_size', type=int, default=64)
  parser.add_argument('--child_eval_batch_size', type=int, default=500)
  parser.add_argument('--child_epochs', type=int, default=150)
  parser.add_argument('--child_layers', type=int, default=3)
  parser.add_argument('--child_nodes', type=int, default=5)
  parser.add_argument('--child_channels', type=int, default=20)
  parser.add_argument('--child_cutout_size', type=int, default=None)
  parser.add_argument('--child_grad_bound', type=float, default=5.0)
  parser.add_argument('--child_lr_max', type=float, default=0.025)
  parser.add_argument('--child_lr_min', type=float, default=0.001)
  parser.add_argument('--child_keep_prob', type=float, default=1.0)
  parser.add_argument('--child_drop_path_keep_prob', type=float, default=0.9)
  parser.add_argument('--child_l2_reg', type=float, default=3e-4)
  parser.add_argument('--child_use_aux_head', action='store_true', default=False)
  parser.add_argument('--child_eval_epochs', type=str, default='30')
  parser.add_argument('--child_arch_pool', type=str, default=None)
  parser.add_argument('--child_lr', type=float, default=0.1)
  parser.add_argument('--child_label_smooth', type=float, default=0.1, help='label smoothing')
  parser.add_argument('--child_gamma', type=float, default=0.97, help='learning rate decay')
  parser.add_argument('--child_decay_period', type=int, default=1, help='epochs between two learning rate decays')
  parser.add_argument('--controller_seed_arch', type=int, default=600)
  parser.add_argument('--controller_expand', type=int, default=None)
  parser.add_argument('--controller_new_arch', type=int, default=300)
  parser.add_argument('--controller_encoder_layers', type=int, default=1)
  parser.add_argument('--controller_encoder_hidden_size', type=int, default=96)
  parser.add_argument('--controller_encoder_emb_size', type=int, default=48)
  parser.add_argument('--controller_mlp_layers', type=int, default=3)
  parser.add_argument('--controller_mlp_hidden_size', type=int, default=200)
  parser.add_argument('--controller_decoder_layers', type=int, default=1)
  parser.add_argument('--controller_decoder_hidden_size', type=int, default=96)
  parser.add_argument('--controller_source_length', type=int, default=40)
  parser.add_argument('--controller_encoder_length', type=int, default=20)
  parser.add_argument('--controller_decoder_length', type=int, default=40)
  parser.add_argument('--controller_encoder_dropout', type=float, default=0)
  parser.add_argument('--controller_mlp_dropout', type=float, default=0.1)
  parser.add_argument('--controller_decoder_dropout', type=float, default=0)
  parser.add_argument('--controller_l2_reg', type=float, default=1e-4)
  parser.add_argument('--controller_encoder_vocab_size', type=int, default=12)
  parser.add_argument('--controller_decoder_vocab_size', type=int, default=12)
  parser.add_argument('--controller_trade_off', type=float, default=0.8)
  parser.add_argument('--controller_epochs', type=int, default=1000)
  parser.add_argument('--controller_batch_size', type=int, default=100)
  parser.add_argument('--controller_lr', type=float, default=0.001)
  parser.add_argument('--controller_optimizer', type=str, default='adam')
  parser.add_argument('--controller_grad_bound', type=float, default=5.0)
  args = parser.parse_args()
  
  utils.create_exp_dir(args.output_dir, scripts_to_save=glob.glob('*.py'))
  
  log_format = '%(asctime)s %(message)s'
  logging.basicConfig(stream=sys.stdout, level=logging.INFO,
      format=log_format, datefmt='%m/%d %I:%M:%S %p')
  
  
  class CrossEntropyLabelSmooth(nn.Module):
  
      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
  
  
  def get_builder(dataset):
      if dataset == 'cifar10':
          return build_cifar10
      elif dataset == 'cifar100':
          return build_cifar100
      else:
          return build_imagenet
      
  
  def build_cifar10(model_state_dict=None, optimizer_state_dict=None, **kwargs):
      epoch = kwargs.pop('epoch')
      ratio = kwargs.pop('ratio')
      train_transform, valid_transform = utils._data_transforms_cifar10(args.child_cutout_size)
      train_data = dset.CIFAR10(root=args.data, train=True, download=True, transform=train_transform)
      valid_data = dset.CIFAR10(root=args.data, train=True, download=True, transform=valid_transform)
      
      num_train = len(train_data)
      assert num_train == len(valid_data)
      indices = list(range(num_train)) 
      split = int(np.floor(ratio * num_train))
      np.random.shuffle(indices)
  
      train_queue = torch.utils.data.DataLoader(
          train_data, batch_size=args.child_batch_size,
          sampler=torch.utils.data.sampler.SubsetRandomSampler(indices[:split]),
          pin_memory=True, num_workers=16)
      valid_queue = torch.utils.data.DataLoader(
          valid_data, batch_size=args.child_eval_batch_size,
          sampler=torch.utils.data.sampler.SubsetRandomSampler(indices[split:num_train]),
          pin_memory=True, num_workers=16)
      
      model = NASWSNetworkCIFAR(10, args.child_layers, args.child_nodes, args.child_channels, args.child_keep_prob, args.child_drop_path_keep_prob,
                         args.child_use_aux_head, args.steps)
      model = model.cuda()
      train_criterion = nn.CrossEntropyLoss().cuda()
      eval_criterion = nn.CrossEntropyLoss().cuda()
      logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
  
      optimizer = torch.optim.SGD(
          model.parameters(),
          args.child_lr_max,
          momentum=0.9,
          weight_decay=args.child_l2_reg,
      )
      if model_state_dict is not None:
          model.load_state_dict(model_state_dict)
      if optimizer_state_dict is not None:
          optimizer.load_state_dict(optimizer_state_dict)
      scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, args.child_epochs, args.child_lr_min, epoch)
      return train_queue, valid_queue, model, train_criterion, eval_criterion, optimizer, scheduler
  
  
  def build_cifar100(model_state_dict=None, optimizer_state_dict=None, **kwargs):
      epoch = kwargs.pop('epoch')
      ratio = kwargs.pop('ratio')
      train_transform, valid_transform = utils._data_transforms_cifar10(args.cutout_size)
      train_data = dset.CIFAR100(root=args.data, train=True, download=True, transform=train_transform)
      valid_data = dset.CIFAR100(root=args.data, train=True, download=True, transform=valid_transform)
  
      num_train = len(train_data)
      assert num_train == len(valid_data)
      indices = list(range(num_train))    
      split = int(np.floor(ratio * num_train))
      np.random.shuffle(indices)
  
      train_queue = torch.utils.data.DataLoader(
          train_data, batch_size=args.child_batch_size,
          sampler=torch.utils.data.sampler.SubsetRandomSampler(indices[:split]),
          pin_memory=True, num_workers=16)
      valid_queue = torch.utils.data.DataLoader(
          valid_data, batch_size=args.child_eval_batch_size,
          sampler=torch.utils.data.sampler.SubsetRandomSampler(indices[split:num_train]),
          pin_memory=True, num_workers=16)
      
      model = NASWSNetworkCIFAR(100, args.child_layers, args.child_nodes, args.child_channels, args.child_keep_prob, args.child_drop_path_keep_prob,
                         args.child_use_aux_head, args.steps)
      model = model.cuda()
      train_criterion = nn.CrossEntropyLoss().cuda()
      eval_criterion = nn.CrossEntropyLoss().cuda()
      logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
  
      optimizer = torch.optim.SGD(
          model.parameters(),
          args.child_lr_max,
          momentum=0.9,
          weight_decay=args.child_l2_reg,
      )
      if model_state_dict is not None:
          model.load_state_dict(model_state_dict)
      if optimizer_state_dict is not None:
          optimizer.load_state_dict(optimizer_state_dict)
      scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, args.child_epochs, args.child_lr_min, epoch)
      return train_queue, valid_queue, model, train_criterion, eval_criterion, optimizer, scheduler
  
  
  def build_imagenet(model_state_dict=None, optimizer_state_dict=None, **kwargs):
      ratio = kwargs.pop('ratio')
      normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
      train_transform = transforms.Compose([
          transforms.RandomResizedCrop(224),
          transforms.RandomHorizontalFlip(),
          transforms.ColorJitter(
              brightness=0.4,
              contrast=0.4,
              saturation=0.4,
              hue=0.2),
          transforms.ToTensor(),
          normalize,
      ])
      valid_transform = transforms.Compose([
          transforms.Resize(256),
          transforms.CenterCrop(224),
          transforms.ToTensor(),
          normalize,
      ])
      if args.zip_file:
          logging.info('Loading data from zip file')
          traindir = os.path.join(args.data, 'train.zip')
          if args.lazy_load:
              train_data = utils.ZipDataset(traindir, train_transform)
          else:
              logging.info('Loading data into memory')
              train_data = utils.InMemoryZipDataset(traindir, train_transform, num_workers=32)
      else:
          logging.info('Loading data from directory')
          traindir = os.path.join(args.data, 'train')
          if args.lazy_load:
              train_data = dset.ImageFolder(traindir, train_transform)
          else:
              logging.info('Loading data into memory')
              train_data = utils.InMemoryDataset(traindir, train_transform, num_workers=32)
         
      num_train = len(train_data)
      indices = list(range(num_train))
      np.random.shuffle(indices)
      split = int(np.floor(ratio * num_train))
      train_indices = sorted(indices[:split])
      valid_indices = sorted(indices[split:])
  
      train_queue = torch.utils.data.DataLoader(
          train_data, batch_size=args.child_batch_size,
          sampler=torch.utils.data.sampler.SubsetRandomSampler(train_indices),
          pin_memory=True, num_workers=16)
      valid_queue = torch.utils.data.DataLoader(
          train_data, batch_size=args.child_eval_batch_size,
          sampler=torch.utils.data.sampler.SubsetRandomSampler(valid_indices),
          pin_memory=True, num_workers=16)
      
      model = NASWSNetworkImageNet(1000, args.child_layers, args.child_nodes, args.child_channels, args.child_keep_prob,
                         args.child_drop_path_keep_prob, args.child_use_aux_head, args.steps)
      model = model.cuda()
      train_criterion = CrossEntropyLabelSmooth(1000, args.child_label_smooth).cuda()
      eval_criterion = nn.CrossEntropyLoss().cuda()
      logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
  
      optimizer = torch.optim.SGD(
          model.parameters(),
          args.child_lr,
          momentum=0.9,
          weight_decay=args.child_l2_reg,
      )
      if model_state_dict is not None:
          model.load_state_dict(model_state_dict)
      if optimizer_state_dict is not None:
          optimizer.load_state_dict(optimizer_state_dict)
      scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.child_decay_period, gamma=args.child_gamma)
      return train_queue, valid_queue, model, train_criterion, eval_criterion, optimizer, scheduler
  
  
  def child_train(train_queue, model, optimizer, global_step, arch_pool, arch_pool_prob, criterion):
      objs = utils.AvgrageMeter()
      top1 = utils.AvgrageMeter()
      top5 = utils.AvgrageMeter()
      model.train()
      for step, (input, target) in enumerate(train_queue):
          input = input.cuda().requires_grad_()
          target = target.cuda()
  
          optimizer.zero_grad()
          # sample an arch to train
          arch = utils.sample_arch(arch_pool, arch_pool_prob)
          logits, aux_logits = model(input, arch, global_step)
          global_step += 1
          loss = criterion(logits, target)
          if aux_logits is not None:
              aux_loss = criterion(aux_logits, target)
              loss += 0.4 * aux_loss
          loss.backward()
          nn.utils.clip_grad_norm_(model.parameters(), args.child_grad_bound)
          optimizer.step()
          
          prec1, prec5 = utils.accuracy(logits, target, topk=(1, 5))
          n = input.size(0)
          objs.update(loss.data, n)
          top1.update(prec1.data, n)
          top5.update(prec5.data, n)
          
          if (step+1) % 100 == 0:
              logging.info('Train %03d loss %e top1 %f top5 %f', step+1, objs.avg, top1.avg, top5.avg)
              logging.info('Arch: %s', ' '.join(map(str, arch[0] + arch[1])))
  
      return top1.avg, objs.avg, global_step
  
  
  def child_valid(valid_queue, model, arch_pool, criterion):
      valid_acc_list = []
      with torch.no_grad():
          model.eval()
          for i, arch in enumerate(arch_pool):
              #for step, (inputs, targets) in enumerate(valid_queue):
              inputs, targets = next(iter(valid_queue))
              inputs = inputs.cuda()
              targets = targets.cuda()
                  
              logits, _ = model(inputs, arch, bn_train=True)
              loss = criterion(logits, targets)
                  
              prec1, prec5 = utils.accuracy(logits, targets, topk=(1, 5))
              valid_acc_list.append(prec1.data/100)
              
              if (i+1) % 100 == 0:
                  logging.info('Valid arch %s\n loss %.2f top1 %f top5 %f', ' '.join(map(str, arch[0] + arch[1])), loss, prec1, prec5)
          
      return valid_acc_list
  
  
  def train_and_evaluate_top_on_cifar10(archs, train_queue, valid_queue):
      res = []
      train_criterion = nn.CrossEntropyLoss().cuda()
      eval_criterion = nn.CrossEntropyLoss().cuda()
      objs = utils.AvgrageMeter()
      top1 = utils.AvgrageMeter()
      top5 = utils.AvgrageMeter()
      for i, arch in enumerate(archs):
          objs.reset()
          top1.reset()
          top5.reset()
          logging.info('Train and evaluate the {} arch'.format(i+1))
          model = NASNetworkCIFAR(args, 10, args.child_layers, args.child_nodes, args.child_channels, 0.6, 0.8,
                          True, args.steps, arch)
          model = model.cuda()
          model.train()
          optimizer = torch.optim.SGD(
              model.parameters(),
              args.child_lr_max,
              momentum=0.9,
              weight_decay=args.child_l2_reg,
          )
          scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, 10, args.child_lr_min)
          global_step = 0
          for e in range(10):
              scheduler.step()
              for step, (input, target) in enumerate(train_queue):
                  input = input.cuda().requires_grad_()
                  target = target.cuda()
  
                  optimizer.zero_grad()
                  # sample an arch to train
                  logits, aux_logits = model(input, global_step)
                  global_step += 1
                  loss = train_criterion(logits, target)
                  if aux_logits is not None:
                      aux_loss = train_criterion(aux_logits, target)
                      loss += 0.4 * aux_loss
                  loss.backward()
                  nn.utils.clip_grad_norm_(model.parameters(), args.child_grad_bound)
                  optimizer.step()
              
                  prec1, prec5 = utils.accuracy(logits, target, topk=(1, 5))
                  n = input.size(0)
                  objs.update(loss.data, n)
                  top1.update(prec1.data, n)
                  top5.update(prec5.data, n)
              
                  if (step+1) % 100 == 0:
                      logging.info('Train epoch %03d %03d loss %e top1 %f top5 %f', e+1, step+1, objs.avg, top1.avg, top5.avg)
          objs.reset()
          top1.reset()
          top5.reset()
          with torch.no_grad():
              model.eval()
              for step, (input, target) in enumerate(valid_queue):
                  input = input.cuda()
                  target = target.cuda()
              
                  logits, _ = model(input)
                  loss = eval_criterion(logits, target)
              
                  prec1, prec5 = utils.accuracy(logits, target, topk=(1, 5))
                  n = input.size(0)
                  objs.update(loss.data, n)
                  top1.update(prec1.data, n)
                  top5.update(prec5.data, n)
              
                  if (step+1) % 100 == 0:
                      logging.info('valid %03d %e %f %f', step+1, objs.avg, top1.avg, top5.avg)
          res.append(top1.avg)
      return res
  
  
  def train_and_evaluate_top_on_cifar100(archs, train_queue, valid_queue):
      res = []
      train_criterion = nn.CrossEntropyLoss().cuda()
      eval_criterion = nn.CrossEntropyLoss().cuda()
      objs = utils.AvgrageMeter()
      top1 = utils.AvgrageMeter()
      top5 = utils.AvgrageMeter()
      for i, arch in enumerate(archs):
          objs.reset()
          top1.reset()
          top5.reset()
          logging.info('Train and evaluate the {} arch'.format(i+1))
          model = NASNetworkCIFAR(args, 100, args.child_layers, args.child_nodes, args.child_channels, 0.6, 0.8,
                          True, args.steps, arch)
          model = model.cuda()
          model.train()
          optimizer = torch.optim.SGD(
              model.parameters(),
              args.child_lr_max,
              momentum=0.9,
              weight_decay=args.child_l2_reg,
          )
          scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, 10, args.child_lr_min)
          global_step = 0
          for e in range(10):
              scheduler.step()
              for step, (input, target) in enumerate(train_queue):
                  input = input.cuda().requires_grad_()
                  target = target.cuda()
  
                  optimizer.zero_grad()
                  # sample an arch to train
                  logits, aux_logits = model(input, global_step)
                  global_step += 1
                  loss = train_criterion(logits, target)
                  if aux_logits is not None:
                      aux_loss = train_criterion(aux_logits, target)
                      loss += 0.4 * aux_loss
                  loss.backward()
                  nn.utils.clip_grad_norm_(model.parameters(), args.child_grad_bound)
                  optimizer.step()
              
                  prec1, prec5 = utils.accuracy(logits, target, topk=(1, 5))
                  n = input.size(0)
                  objs.update(loss.data, n)
                  top1.update(prec1.data, n)
                  top5.update(prec5.data, n)
              
                  if (step+1) % 100 == 0:
                      logging.info('Train %3d %03d loss %e top1 %f top5 %f', e+1, step+1, objs.avg, top1.avg, top5.avg)
          objs.reset()
          top1.reset()
          top5.reset()
          with torch.no_grad():
              model.eval()
              for step, (input, target) in enumerate(valid_queue):
                  input = input.cuda()
                  target = target.cuda()
              
                  logits, _ = model(input)
                  loss = eval_criterion(logits, target)
              
                  prec1, prec5 = utils.accuracy(logits, target, topk=(1, 5))
                  n = input.size(0)
                  objs.update(loss.data, n)
                  top1.update(prec1.data, n)
                  top5.update(prec5.data, n)
              
                  if (step+1) % 100 == 0:
                      logging.info('valid %03d %e %f %f', step+1, objs.avg, top1.avg, top5.avg)
          res.append(top1.avg)
      return res
  
  
  def train_and_evaluate_top_on_imagenet(archs, train_queue, valid_queue):
      res = []
      train_criterion = nn.CrossEntropyLoss().cuda()
      eval_criterion = nn.CrossEntropyLoss().cuda()
      objs = utils.AvgrageMeter()
      top1 = utils.AvgrageMeter()
      top5 = utils.AvgrageMeter()
      for i, arch in enumerate(archs):
          objs.reset()
          top1.reset()
          top5.reset()
          logging.info('Train and evaluate the {} arch'.format(i+1))
          model = NASNetworkImageNet(args, 1000, args.child_layers, args.child_nodes, args.child_channels, 1.0, 1.0,
                          True, args.steps, arch)
          model = model.cuda()
          model.train()
          optimizer = torch.optim.SGD(
              model.parameters(),
              args.child_lr,
              momentum=0.9,
              weight_decay=args.child_l2_reg,
          )
          for step, (input, target) in enumerate(train_queue):
              input = input.cuda().requires_grad_()
              target = target.cuda()
  
              optimizer.zero_grad()
              # sample an arch to train
              logits, aux_logits = model(input, step)
              loss = train_criterion(logits, target)
              if aux_logits is not None:
                  aux_loss = train_criterion(aux_logits, target)
                  loss += 0.4 * aux_loss
              loss.backward()
              nn.utils.clip_grad_norm_(model.parameters(), args.child_grad_bound)
              optimizer.step()
              
              prec1, prec5 = utils.accuracy(logits, target, topk=(1, 5))
              n = input.size(0)
              objs.update(loss.data, n)
              top1.update(prec1.data, n)
              top5.update(prec5.data, n)
              
              if (step+1) % 100 == 0:
                  logging.info('Train %03d loss %e top1 %f top5 %f', step+1, objs.avg, top1.avg, top5.avg)
              if step+1 == 500:
                  break
  
          objs.reset()
          top1.reset()
          top5.reset()
          with torch.no_grad():
              model.eval()
              for step, (input, target) in enumerate(valid_queue):
                  input = input.cuda()
                  target = target.cuda()
              
                  logits, _ = model(input)
                  loss = eval_criterion(logits, target)
              
                  prec1, prec5 = utils.accuracy(logits, target, topk=(1, 5))
                  n = input.size(0)
                  objs.update(loss.data, n)
                  top1.update(prec1.data, n)
                  top5.update(prec5.data, n)
              
                  if (step+1) % 100 == 0:
                      logging.info('valid %03d %e %f %f', step+1, objs.avg, top1.avg, top5.avg)
          res.append(top1.avg)
      return res
          
          
  def nao_train(train_queue, model, optimizer):
      objs = utils.AvgrageMeter()
      mse = utils.AvgrageMeter()
      nll = utils.AvgrageMeter()
      model.train()
      for step, sample in enumerate(train_queue):
          encoder_input = sample['encoder_input']
          encoder_target = sample['encoder_target']
          decoder_input = sample['decoder_input']
          decoder_target = sample['decoder_target']
          
          encoder_input = encoder_input.cuda()
          encoder_target = encoder_target.cuda().requires_grad_()
          decoder_input = decoder_input.cuda()
          decoder_target = decoder_target.cuda()
          
          optimizer.zero_grad()
          predict_value, log_prob, arch = model(encoder_input, decoder_input)
          loss_1 = F.mse_loss(predict_value.squeeze(), encoder_target.squeeze())
          loss_2 = F.nll_loss(log_prob.contiguous().view(-1, log_prob.size(-1)), decoder_target.view(-1))
          loss = args.controller_trade_off * loss_1 + (1 - args.controller_trade_off) * loss_2
          loss.backward()
          torch.nn.utils.clip_grad_norm_(model.parameters(), args.controller_grad_bound)
          optimizer.step()
          
          n = encoder_input.size(0)
          objs.update(loss.data, n)
          mse.update(loss_1.data, n)
          nll.update(loss_2.data, n)
          
      return objs.avg, mse.avg, nll.avg
  
  
  def nao_valid(queue, model):
      inputs = []
      targets = []
      predictions = []
      archs = []
      with torch.no_grad():
          model.eval()
          for step, sample in enumerate(queue):
              encoder_input = sample['encoder_input']
              encoder_target = sample['encoder_target']
              decoder_target = sample['decoder_target']
              
              encoder_input = encoder_input.cuda()
              encoder_target = encoder_target.cuda()
              decoder_target = decoder_target.cuda()
              
              predict_value, logits, arch = model(encoder_input)
              n = encoder_input.size(0)
              inputs += encoder_input.data.squeeze().tolist()
              targets += encoder_target.data.squeeze().tolist()
              predictions += predict_value.data.squeeze().tolist()
              archs += arch.data.squeeze().tolist()
      pa = utils.pairwise_accuracy(targets, predictions)
      hd = utils.hamming_distance(inputs, archs)
      return pa, hd
  
  
  def nao_infer(queue, model, step, direction='+'):
      new_arch_list = []
      model.eval()
      for i, sample in enumerate(queue):
          encoder_input = sample['encoder_input']
          encoder_input = encoder_input.cuda()
          model.zero_grad()
          new_arch = model.generate_new_arch(encoder_input, step, direction=direction)
          new_arch_list.extend(new_arch.data.squeeze().tolist())
      return new_arch_list
  
  
  def main():
      if not torch.cuda.is_available():
          logging.info('no gpu device available')
          sys.exit(1)
          
      random.seed(args.seed)
      np.random.seed(args.seed)
      torch.manual_seed(args.seed)
      torch.cuda.manual_seed(args.seed)
      torch.cuda.manual_seed_all(args.seed)
      cudnn.enabled = True
      cudnn.benchmark = True
      cudnn.deterministic = True
      
      args.steps = int(np.ceil(45000 / args.child_batch_size)) * args.child_epochs
  
      logging.info("args = %s", args)
  
      if args.child_arch_pool is not None:
          logging.info('Architecture pool is provided, loading')
          with open(args.child_arch_pool) as f:
              archs = f.read().splitlines()
              archs = list(map(utils.build_dag, archs))
              child_arch_pool = archs
      elif os.path.exists(os.path.join(args.output_dir, 'arch_pool')):
          logging.info('Architecture pool is founded, loading')
          with open(os.path.join(args.output_dir, 'arch_pool')) as f:
              archs = f.read().splitlines()
              archs = list(map(utils.build_dag, archs))
              child_arch_pool = archs
      else:
          child_arch_pool = None
  
      child_eval_epochs = eval(args.child_eval_epochs)
      build_fn = get_builder(args.dataset)
      train_queue, valid_queue, model, train_criterion, eval_criterion, optimizer, scheduler = build_fn(ratio=0.9, epoch=-1)
  
      nao = NAO(
          args.controller_encoder_layers,
          args.controller_encoder_vocab_size,
          args.controller_encoder_hidden_size,
          args.controller_encoder_dropout,
          args.controller_encoder_length,
          args.controller_source_length,
          args.controller_encoder_emb_size,
          args.controller_mlp_layers,
          args.controller_mlp_hidden_size,
          args.controller_mlp_dropout,
          args.controller_decoder_layers,
          args.controller_decoder_vocab_size,
          args.controller_decoder_hidden_size,
          args.controller_decoder_dropout,
          args.controller_decoder_length,
      )
      nao = nao.cuda()
      logging.info("Encoder-Predictor-Decoder param size = %fMB", utils.count_parameters_in_MB(nao))
  
      # Train child model
      if child_arch_pool is None:
          logging.info('Architecture pool is not provided, randomly generating now')
          child_arch_pool = utils.generate_arch(args.controller_seed_arch, args.child_nodes, 5)  # [[[conv],[reduc]]]
      child_arch_pool_prob = None
  
      eval_points = utils.generate_eval_points(child_eval_epochs, 0, args.child_epochs)
      step = 0
      for epoch in range(1, args.child_epochs + 1):
          scheduler.step()
          lr = scheduler.get_lr()[0]
          logging.info('epoch %d lr %e', epoch, lr)
          # sample an arch to train
          train_acc, train_obj, step = child_train(train_queue, model, optimizer, step, child_arch_pool, child_arch_pool_prob, train_criterion)
          logging.info('train_acc %f', train_acc)
      
          if epoch not in eval_points:
              continue
          # Evaluate seed archs
          valid_accuracy_list = child_valid(valid_queue, model, child_arch_pool, eval_criterion)
  
          # Output archs and evaluated error rate
          old_archs = child_arch_pool
          old_archs_perf = valid_accuracy_list
  
          old_archs_sorted_indices = np.argsort(old_archs_perf)[::-1]
          old_archs = [old_archs[i] for i in old_archs_sorted_indices]
          old_archs_perf = [old_archs_perf[i] for i in old_archs_sorted_indices]
          with open(os.path.join(args.output_dir, 'arch_pool.{}'.format(epoch)), 'w') as fa:
              with open(os.path.join(args.output_dir, 'arch_pool.perf.{}'.format(epoch)), 'w') as fp:
                  with open(os.path.join(args.output_dir, 'arch_pool'), 'w') as fa_latest:
                      with open(os.path.join(args.output_dir, 'arch_pool.perf'), 'w') as fp_latest:
                          for arch, perf in zip(old_archs, old_archs_perf):
                              arch = ' '.join(map(str, arch[0] + arch[1]))
                              fa.write('{}\n'.format(arch))
                              fa_latest.write('{}\n'.format(arch))
                              fp.write('{}\n'.format(perf))
                              fp_latest.write('{}\n'.format(perf))
                              
          if epoch == args.child_epochs:
              break
  
          # Train Encoder-Predictor-Decoder
          logging.info('Training Encoder-Predictor-Decoder')
          encoder_input = list(map(lambda x: utils.parse_arch_to_seq(x[0], 2) + utils.parse_arch_to_seq(x[1], 2), old_archs))
          # [[conv, reduc]]
          min_val = min(old_archs_perf)
          max_val = max(old_archs_perf)
          encoder_target = [(i - min_val) / (max_val - min_val) for i in old_archs_perf]
  
          if args.controller_expand:
              dataset = list(zip(encoder_input, encoder_target))
              n = len(dataset)
              ratio = 0.9
              split = int(n*ratio)
              np.random.shuffle(dataset)
              encoder_input, encoder_target = list(zip(*dataset))
              train_encoder_input = list(encoder_input[:split])
              train_encoder_target = list(encoder_target[:split])
              valid_encoder_input = list(encoder_input[split:])
              valid_encoder_target = list(encoder_target[split:])
              for _ in range(args.controller_expand-1):
                  for src, tgt in zip(encoder_input[:split], encoder_target[:split]):
                      a = np.random.randint(0, args.child_nodes)
                      b = np.random.randint(0, args.child_nodes)
                      src = src[:4 * a] + src[4 * a + 2:4 * a + 4] + \
                              src[4 * a:4 * a + 2] + src[4 * (a + 1):20 + 4 * b] + \
                              src[20 + 4 * b + 2:20 + 4 * b + 4] + src[20 + 4 * b:20 + 4 * b + 2] + \
                              src[20 + 4 * (b + 1):]
                      train_encoder_input.append(src)
                      train_encoder_target.append(tgt)
          else:
              train_encoder_input = encoder_input
              train_encoder_target = encoder_target
              valid_encoder_input = encoder_input
              valid_encoder_target = encoder_target
          logging.info('Train data: {}\tValid data: {}'.format(len(train_encoder_input), len(valid_encoder_input)))
  
          nao_train_dataset = utils.NAODataset(train_encoder_input, train_encoder_target, True, swap=True if args.controller_expand is None else False)
          nao_valid_dataset = utils.NAODataset(valid_encoder_input, valid_encoder_target, False)
          nao_train_queue = torch.utils.data.DataLoader(
              nao_train_dataset, batch_size=args.controller_batch_size, shuffle=True, pin_memory=True)
          nao_valid_queue = torch.utils.data.DataLoader(
              nao_valid_dataset, batch_size=args.controller_batch_size, shuffle=False, pin_memory=True)
          nao_optimizer = torch.optim.Adam(nao.parameters(), lr=args.controller_lr, weight_decay=args.controller_l2_reg)
          for nao_epoch in range(1, args.controller_epochs+1):
              nao_loss, nao_mse, nao_ce = nao_train(nao_train_queue, nao, nao_optimizer)
              logging.info("epoch %04d train loss %.6f mse %.6f ce %.6f", nao_epoch, nao_loss, nao_mse, nao_ce)
              if nao_epoch % 100 == 0:
                  pa, hs = nao_valid(nao_valid_queue, nao)
                  logging.info("Evaluation on valid data")
                  logging.info('epoch %04d pairwise accuracy %.6f hamming distance %.6f', epoch, pa, hs)
  
          # Generate new archs
          new_archs = []
          max_step_size = 50
          predict_step_size = 0
          top100_archs = list(map(lambda x: utils.parse_arch_to_seq(x[0], 2) + utils.parse_arch_to_seq(x[1], 2), old_archs[:100]))
          nao_infer_dataset = utils.NAODataset(top100_archs, None, False)
          nao_infer_queue = torch.utils.data.DataLoader(
              nao_infer_dataset, batch_size=len(nao_infer_dataset), shuffle=False, pin_memory=True)
          while len(new_archs) < args.controller_new_arch:
              predict_step_size += 1
              logging.info('Generate new architectures with step size %d', predict_step_size)
              new_arch = nao_infer(nao_infer_queue, nao, predict_step_size, direction='+')
              for arch in new_arch:
                  if arch not in encoder_input and arch not in new_archs:
                      new_archs.append(arch)
                  if len(new_archs) >= args.controller_new_arch:
                      break
              logging.info('%d new archs generated now', len(new_archs))
              if predict_step_size > max_step_size:
                  break
                  # [[conv, reduc]]
          new_archs = list(map(lambda x: utils.parse_seq_to_arch(x, 2), new_archs))  # [[[conv],[reduc]]]
          num_new_archs = len(new_archs)
          logging.info("Generate %d new archs", num_new_archs)
          # replace bottom archs
          new_arch_pool = old_archs[:len(old_archs) - num_new_archs] + new_archs
          logging.info("Totally %d architectures now to train", len(new_arch_pool))
  
          child_arch_pool = new_arch_pool
          with open(os.path.join(args.output_dir, 'arch_pool'), 'w') as f:
              for arch in new_arch_pool:
                  arch = ' '.join(map(str, arch[0] + arch[1]))
                  f.write('{}\n'.format(arch))
  
          
          child_arch_pool_prob = None
  
      logging.info('Finish Searching')
      logging.info('Reranking top 5 architectures')
      # reranking top 5
      top_archs = old_archs[:5]
      if args.dataset == 'cifar10':
          top_archs_perf = train_and_evaluate_top_on_cifar10(top_archs, train_queue, valid_queue)
      elif args.dataset == 'cifar100':
          top_archs_perf = train_and_evaluate_top_on_cifar100(top_archs, train_queue, valid_queue)
      else:
          top_archs_perf = train_and_evaluate_top_on_imagenet(top_archs, train_queue, valid_queue)
      top_archs_sorted_indices = np.argsort(top_archs_perf)[::-1]
      top_archs = [top_archs[i] for i in top_archs_sorted_indices]
      top_archs_perf = [top_archs_perf[i] for i in top_archs_sorted_indices]
      with open(os.path.join(args.output_dir, 'arch_pool.final'), 'w') as fa:
          with open(os.path.join(args.output_dir, 'arch_pool.perf.final'), 'w') as fp:
              for arch, perf in zip(top_archs, top_archs_perf):
                  arch = ' '.join(map(str, arch[0] + arch[1]))
                  fa.write('{}\n'.format(arch))
                  fp.write('{}\n'.format(perf))
    
  
  if __name__ == '__main__':
      main()
  