code: def main(args):
      """The main function for supernet pre-training and subnet fine-tuning."""
      logging.basicConfig(
          format="[%(asctime)s] [p%(process)s] [%(pathname)s\
              :%(lineno)d] [%(levelname)s] %(message)s",
          level=logging.INFO,
          handlers=[
              logging.FileHandler(args.log_file, mode="w"),
              logging.StreamHandler(),
          ],
      )
  
      # print the information of arguments
      for arg in vars(args):
          s = arg + ": " + str(getattr(args, arg))
          logging.info(s)
  
      # for 106 landmarks
      num_points = 106
      # list of device ids, and the number of workers for data loading
      device_ids = [int(id) for id in args.dev_id.split(",")]
      dev_num = len(device_ids)
      num_workers = 4 * dev_num
  
      # import subnet for fine-tuning
      from lib.subnet import PFLDInference, AuxiliaryNet
  
      # the configuration for training control
      nas_config = NASConfig(
          model_dir=args.snapshot,
          search_space=search_space,
      )
      # look-up table with information of search space, flops per block, etc.
      lookup_table = LookUpTable(
          config=nas_config, primitives=PRIMITIVES
      )
  
      check = torch.load(
          args.supernet, map_location=torch.device("cpu")
      )
      sampled_arch = check["arch_sample"]
      logging.info(sampled_arch)
      # create subnet
      pfld_backbone = PFLDInference(
          lookup_table, sampled_arch, num_points
      )
  
      # pre-load the weights from pre-trained supernet
      state_dict = check["pfld_backbone"]
      supernet_sample(
          pfld_backbone, state_dict, sampled_arch, lookup_table
      )
  
      # the auxiliary-net of PFLD to predict the pose angle
      auxiliarynet = AuxiliaryNet()
  
      # DataParallel
      pfld_backbone = torch.nn.DataParallel(
          pfld_backbone, device_ids=device_ids
      )
      pfld_backbone.to(device)
      auxiliarynet = torch.nn.DataParallel(
          auxiliarynet, device_ids=device_ids
      )
      auxiliarynet.to(device)
  
      # main task loss
      criterion = PFLDLoss()
  
      # optimizer / scheduler for weight train
      optimizer = torch.optim.RMSprop(
          [
              {"params": pfld_backbone.parameters()},
              {"params": auxiliarynet.parameters()},
          ],
          lr=args.base_lr,
          momentum=0.0,
          weight_decay=args.weight_decay,
      )
  
      scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
          optimizer, T_max=args.end_epoch, last_epoch=-1
      )
  
      # data argmentation and dataloader
      transform = torchvision.transforms.Compose(
          [torchvision.transforms.ToTensor()]
      )
      # the landmark dataset with 106 points is default used
      train_dataset = PFLDDatasets(
          os.path.join(args.data_root, "train_data/list.txt"),
          transform,
          data_root=args.data_root,
          img_size=args.img_size,
      )
      dataloader = DataLoader(
          train_dataset,
          batch_size=args.train_batchsize,
          shuffle=True,
          num_workers=num_workers,
          pin_memory=True,
          drop_last=False,
      )
  
      val_dataset = PFLDDatasets(
          os.path.join(args.data_root, "test_data/list.txt"),
          transform,
          data_root=args.data_root,
          img_size=args.img_size,
      )
      val_dataloader = DataLoader(
          val_dataset,
          batch_size=args.val_batchsize,
          shuffle=False,
          num_workers=num_workers,
          pin_memory=True,
      )
  
      # start finetune
      ckpt_path = args.snapshot
      val_nme = 1e6
  
      for epoch in range(0, args.end_epoch):
          logging.info("\n--------Train epoch: %d--------\n", epoch + 1)
          # update the weight parameters
          train_epoch(
              pfld_backbone,
              auxiliarynet,
              criterion,
              dataloader,
              device,
              epoch,
              optimizer,
              logging,
          )
          # adjust learning rate
          scheduler.step()
  
          # validate
          _, nme = validate(
              pfld_backbone,
              auxiliarynet,
              val_dataloader,
              device,
              logging,
          )
  
          if epoch % 10 == 0:
              filename = os.path.join(
                  ckpt_path, "checkpoint_%s.pth" % epoch
              )
              save_checkpoint(
                  pfld_backbone,
                  auxiliarynet,
                  optimizer,
                  filename,
                  logging,
              )
  
          if nme < val_nme:
              filename = os.path.join(ckpt_path, "checkpoint_best.pth")
              save_checkpoint(
                  pfld_backbone,
                  auxiliarynet,
                  optimizer,
                  filename,
                  logging,
              )
              val_nme = nme
          logging.info("Best nme: {:.4f}".format(val_nme))
  
file_path: nni/examples/nas/legacy/oneshot/pfld/retrain.py
imports: from __future__ import absolute_import, division, print_function
  import argparse
  import logging
  import os
  import time
  import torch
  import torchvision
  import numpy
  from datasets import PFLDDatasets
  from lib.builder import search_space
  from lib.ops import PRIMITIVES
  from lib.utils import PFLDLoss, accuracy
  from nni.algorithms.nas.pytorch.fbnet import LookUpTable, NASConfig, supernet_sample
  from nni.nas.pytorch.utils import AverageMeter
  from torch.utils.data import DataLoader
label_desc: Use image node_id 4 for calling a global function with example usage: main(args) without return types