text: """ Search cell """
  import os
  import torch
  import torch.nn as nn
  import numpy as np
  import time
  from tensorboardX import SummaryWriter
  from config import SearchConfig
  import utils
  from models.search_cnn import SearchCNNController
  from architect import Architect
  from visualize import plot
  from torchsampler import ImbalancedDatasetSampler
  
  config = SearchConfig()
  
  device = torch.device("cuda")
  
  # tensorboard
  writer = SummaryWriter(log_dir=os.path.join(config.path, "tb"))
  writer.add_text('config', config.as_markdown(), 0)
  
  logger = utils.get_logger(os.path.join(config.path, "{}.log".format(config.name)))
  config.print_params(logger.info)
  
  
  def main():
      logger.info("Logger is set - training start")
  
      # set default gpu device id
      torch.cuda.set_device(config.gpus[0])
  
      # set seed
      np.random.seed(config.seed)
      torch.manual_seed(config.seed)
      torch.cuda.manual_seed_all(config.seed)
  
      torch.backends.cudnn.benchmark = True
  
      # get data with meta info
      input_size, input_channels, n_classes, train_data, valid_data = utils.get_data(
          config.dataset, config.data_path, cutout_length=config.cutout_length, validation=True, val_path=config.val_set, fealen=config.fealen, inputsize=config.input_size)
  
      net_crit = nn.CrossEntropyLoss().to(device)
      model = SearchCNNController(input_channels, config.init_channels, n_classes, config.layers,
                                  net_crit, device_ids=config.gpus)
      model = model.to(device)
  
      # weights optimizer
      w_optim = torch.optim.SGD(model.weights(), config.w_lr, momentum=config.w_momentum,
                                weight_decay=config.w_weight_decay)
      # alphas optimizer
      alpha_optim = torch.optim.Adam(model.alphas(), config.alpha_lr, betas=(0.5, 0.999),
                                     weight_decay=config.alpha_weight_decay)
  
      # split data to train/validation
  
      # n_train = len(train_data)
      # split = n_train // 2
      # indices = list(range(n_train))
      # train_sampler = torch.utils.data.sampler.SubsetRandomSampler(indices[:split])
      # valid_sampler = torch.utils.data.sampler.SubsetRandomSampler(indices[split:])
      # train_loader = torch.utils.data.DataLoader(train_data,
      #                                            batch_size=config.batch_size,
      #                                            sampler=train_sampler,
      #                                            num_workers=config.workers,
      #                                            pin_memory=True)
      # valid_loader = torch.utils.data.DataLoader(train_data,
      #                                            batch_size=config.batch_size,
      #                                            sampler=valid_sampler,
      #                                            num_workers=config.workers,
      #                                            pin_memory=True)
      
      train_loader = torch.utils.data.DataLoader(train_data,
                                                 batch_size=config.batch_size,
                                              #    shuffle=True,
                                                 sampler=ImbalancedDatasetSampler(train_data),
                                                 num_workers=config.workers,
                                                 pin_memory=True,
                                                 drop_last=True)
      valid_loader = torch.utils.data.DataLoader(valid_data,
                                                 batch_size=config.batch_size,
                                                 shuffle=False,
                                                 num_workers=config.workers,
                                                 pin_memory=True,
                                                 drop_last=True)
      
      lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
          w_optim, config.epochs, eta_min=config.w_lr_min)
      architect = Architect(model, config.w_momentum, config.w_weight_decay)
  
      # training loop
      best_top1 = 0.
      for epoch in range(config.epochs):
          lr_scheduler.step()
          lr = lr_scheduler.get_lr()[0]
  
          model.print_alphas(logger)
  
          # training
          train(train_loader, valid_loader, model, architect, w_optim, alpha_optim, lr, epoch)
  
          # validation
          cur_step = (epoch+1) * len(train_loader)
          top1 = validate(valid_loader, model, epoch, cur_step)
  
          # log
          # genotype
          genotype = model.genotype()
          logger.info("genotype = {}".format(genotype))
  
          # genotype as a image
          plot_path = os.path.join(config.plot_path, "EP{:02d}".format(epoch+1))
          caption = "Epoch {}".format(epoch+1)
          plot(genotype.normal, plot_path + "-normal", caption)
          plot(genotype.reduce, plot_path + "-reduce", caption)
  
          # save
          if best_top1 < top1:
              best_top1 = top1
              best_genotype = genotype
              is_best = True
          else:
              is_best = False
          utils.save_checkpoint(model, config.path, is_best)
          print("")
  
      logger.info("Final best Prec@1 = {:.4%}".format(best_top1))
      logger.info("Best Genotype = {}".format(best_genotype))
  
  
  def train(train_loader, valid_loader, model, architect, w_optim, alpha_optim, lr, epoch):
      confusion = utils.SumMeter()
      losses = utils.AverageMeter()
  
      cur_step = epoch*len(train_loader)
      writer.add_scalar('train/lr', lr, cur_step)
  
      model.train()
  
      for step, ((trn_X, trn_y), (val_X, val_y)) in enumerate(zip(train_loader, valid_loader)):
          trn_X, trn_y = trn_X.to(device, non_blocking=True), trn_y.to(device, non_blocking=True)
          val_X, val_y = val_X.to(device, non_blocking=True), val_y.to(device, non_blocking=True)
          N = trn_X.size(0)
  
          # phase 2. architect step (alpha)
          alpha_optim.zero_grad()
          architect.unrolled_backward(trn_X, trn_y, val_X, val_y, lr, w_optim)
          alpha_optim.step()
  
          # phase 1. child network step (w)
          w_optim.zero_grad()
          logits = model(trn_X)
          loss = model.criterion(logits, trn_y)
          loss.backward()
          # gradient clipping
          nn.utils.clip_grad_norm_(model.weights(), config.w_grad_clip)
          w_optim.step()
  
          conf_mat = utils.binary_conf_mat(logits, trn_y)
          losses.update(loss.item(), N)
          confusion.update(conf_mat)
          if conf_mat[1, 0] + conf_mat[1, 1] > 0:
              recall = conf_mat[1, 1] / (conf_mat[1, 0] + conf_mat[1, 1])
          else:
              recall = 0
          fa = conf_mat[0, 1]
          f1_score = utils.binary_f1_score(conf_mat)
  
          if step % config.print_freq == 0 or step == len(train_loader)-1:
              logger.info(
                  "Train: [{:3d}/{}] Step {:03d}/{:03d} Loss: {:.3f} Recall: {:.3f} FA: {:d} F1: {:3f}".format(
                      epoch+1, config.epochs, step, len(train_loader)-1, losses.avg, recall, fa, f1_score))
              logger.info(conf_mat.flatten())
  
          writer.add_scalar('train/loss', loss.item(), cur_step)
          # writer.add_scalar('train/top1', prec1.item(), cur_step)
          # writer.add_scalar('train/top5', prec5.item(), cur_step)
          cur_step += 1
  
      recall = confusion.val[1,1]/(confusion.val[1,0]+confusion.val[1,1])
      fa = confusion.val[0,1]
      f1_score = utils.binary_f1_score(confusion.val)
      logger.info("Train: [{:3d}/{}] Final Recall {:.4%} FA {:d} F1: {:.4%}".format(epoch+1, config.epochs, recall, fa, f1_score))
      logger.info(confusion.val.flatten())
  
  
  def validate(valid_loader, model, epoch, cur_step):
      confusion = utils.SumMeter()
      losses = utils.AverageMeter()
  
      model.eval()
  
      with torch.no_grad():
          start = time.time()
          for step, (X, y) in enumerate(valid_loader):
              X, y = X.to(device, non_blocking=True), y.to(device, non_blocking=True)
              N = X.size(0)
  
              logits = model(X)
              loss = model.criterion(logits, y)
  
              conf_mat = utils.binary_conf_mat(logits, y)
              losses.update(loss.item(), N)
              confusion.update(conf_mat)
              if conf_mat[1, 0] + conf_mat[1, 1] > 0:
                  recall = conf_mat[1, 1] / (conf_mat[1, 0] + conf_mat[1, 1])
              else:
                  recall = 0
              fa = conf_mat[0, 1]
              f1_score = utils.binary_f1_score(conf_mat)
  
              if step % config.print_freq == 0 or step == len(valid_loader)-1:
                  logger.info(
                      "Valid: [{:3d}/{}] Step {:03d}/{:03d} Loss: {:.3f} Recall: {:.3f} FA: {:d} F1: {:3f}".format(
                          epoch+1, config.epochs, step, len(valid_loader)-1, losses.avg, recall, fa, f1_score))
                  logger.info(conf_mat.flatten())
      val_time = time.time() - start
      recall = confusion.val[1,1]/(confusion.val[1,0]+confusion.val[1,1])
      fa = confusion.val[0,1]
      f1_score = utils.binary_f1_score(confusion.val)
  
      writer.add_scalar('val/loss', losses.avg, epoch)
      writer.add_scalar('val/recall', recall, epoch)
      writer.add_scalar('val/fa', fa, epoch)
      writer.add_scalar('val/f1_score', f1_score, epoch)
      
      logger.info("Valid: [{:3d}/{}] Final Recall {:.4%} FA {:d} F1: {:.4%}".format(epoch+1, config.epochs, recall, fa, f1_score))
      logger.info(confusion.val.flatten())
      logger.info("Valid: Time Spent: %fs"%val_time)
  
      return f1_score
  
  
  if __name__ == "__main__":
      main()
  