code: def main():
      config = RetrainConfig()
      main_proc = not config.distributed or config.local_rank == 0
      if config.distributed:
          torch.cuda.set_device(config.local_rank)
          torch.distributed.init_process_group(
              backend="nccl",
              init_method=config.dist_url,
              rank=config.local_rank,
              world_size=config.world_size,
          )
      if main_proc:
          os.makedirs(config.output_path, exist_ok=True)
      if config.distributed:
          torch.distributed.barrier()
      logger = utils.get_logger(
          os.path.join(config.output_path, "search.log")
      )
      if main_proc:
          config.print_params(logger.info)
      utils.reset_seed(config.seed)
  
      loaders, samplers = get_augment_datasets(config)
      train_loader, valid_loader = loaders
      train_sampler, valid_sampler = samplers
  
      model = Model(
          config.dataset,
          config.layers,
          in_channels=config.input_channels,
          channels=config.init_channels,
          retrain=True,
      ).cuda()
      if config.label_smooth > 0:
          criterion = utils.CrossEntropyLabelSmooth(
              config.n_classes, config.label_smooth
          )
      else:
          criterion = nn.CrossEntropyLoss()
  
      fixed_arc_path = os.path.join(
          config.output_path, config.arc_checkpoint
      )
      with open(fixed_arc_path, "r") as f:
          fixed_arc = json.load(f)
      fixed_arc = utils.encode_tensor(fixed_arc, torch.device("cuda"))
      genotypes = utils.parse_results(fixed_arc, n_nodes=4)
      genotypes_dict = {i: genotypes for i in range(3)}
      apply_fixed_architecture(model, fixed_arc_path)
      param_size = utils.param_size(
          model,
          criterion,
          [3, 32, 32] if "cifar" in config.dataset else [3, 224, 224],
      )
  
      if main_proc:
          logger.info("Param size: %.6f", param_size)
          logger.info("Genotype: %s", genotypes)
  
      # change training hyper parameters according to cell type
      if "cifar" in config.dataset:
          if param_size < 3.0:
              config.weight_decay = 3e-4
              config.drop_path_prob = 0.2
          elif 3.0 < param_size < 3.5:
              config.weight_decay = 3e-4
              config.drop_path_prob = 0.3
          else:
              config.weight_decay = 5e-4
              config.drop_path_prob = 0.3
  
      if config.distributed:
          apex.parallel.convert_syncbn_model(model)
          model = DistributedDataParallel(model, delay_allreduce=True)
  
      optimizer = torch.optim.SGD(
          model.parameters(),
          config.lr,
          momentum=config.momentum,
          weight_decay=config.weight_decay,
      )
      lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
          optimizer, config.epochs, eta_min=1e-6
      )
  
      best_top1 = best_top5 = 0.0
      for epoch in range(config.epochs):
          drop_prob = config.drop_path_prob * epoch / config.epochs
          if config.distributed:
              model.module.drop_path_prob(drop_prob)
          else:
              model.drop_path_prob(drop_prob)
          # training
          if config.distributed:
              train_sampler.set_epoch(epoch)
          train(
              logger,
              config,
              train_loader,
              model,
              optimizer,
              criterion,
              epoch,
              main_proc,
          )
  
          # validation
          top1, top5 = validate(
              logger,
              config,
              valid_loader,
              model,
              criterion,
              epoch,
              main_proc,
          )
          best_top1 = max(best_top1, top1)
          best_top5 = max(best_top5, top5)
          lr_scheduler.step()
  
      logger.info(
          "Final best Prec@1 = %.4f Prec@5 = %.4f", best_top1, best_top5
      )
  
file_path: nni/examples/nas/legacy/cdarts/retrain.py
imports: import json
  import logging
  import os
  import time
  from argparse import ArgumentParser
  import torch
  import torch.nn
  import apex
  import datasets
  import utils
  from apex.parallel import DistributedDataParallel
  from config import RetrainConfig
  from datasets.cifar import get_augment_datasets
  from model import Model
  from nni.nas.pytorch.fixed import apply_fixed_architecture
  from nni.nas.pytorch.utils import AverageMeterGroup
label_desc: Use image node_id 3 for calling a global function with example usage: main() without return types