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
  
      # random seed
      manual_seed = 1
      np.random.seed(manual_seed)
      torch.manual_seed(manual_seed)
      torch.cuda.manual_seed_all(manual_seed)
  
      # import supernet for block-wise DNAS pre-training
      from lib.supernet import PFLDInference, AuxiliaryNet
  
      # the configuration for training control
      nas_config = NASConfig(
          perf_metric=args.perf_metric,
          lut_load=args.lut_load,
          model_dir=args.snapshot,
          nas_lr=args.theta_lr,
          mode=args.mode,
          alpha=args.alpha,
          beta=args.beta,
          search_space=search_space,
          start_epoch=args.start_epoch,
      )
      # look-up table with information of search space, flops per block, etc.
      lookup_table = LookUpTable(
          config=nas_config, primitives=PRIMITIVES
      )
  
      # create supernet
      pfld_backbone = PFLDInference(lookup_table, num_points)
      # the auxiliary-net of PFLD to predict the pose angle
      auxiliarynet = AuxiliaryNet()
  
      # main task loss
      criterion = PFLDLoss()
  
      # optimizer for weight train
      if args.opt == "adam":
          optimizer = torch.optim.AdamW(
              [
                  {"params": pfld_backbone.parameters()},
                  {"params": auxiliarynet.parameters()},
              ],
              lr=args.base_lr,
              weight_decay=args.weight_decay,
          )
      elif args.opt == "rms":
          optimizer = torch.optim.RMSprop(
              [
                  {"params": pfld_backbone.parameters()},
                  {"params": auxiliarynet.parameters()},
              ],
              lr=args.base_lr,
              momentum=0.0,
              weight_decay=args.weight_decay,
          )
  
      # 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,
      )
  
      # create the trainer, then search/finetune
      trainer = PFLDTrainer(
          pfld_backbone,
          auxiliarynet,
          optimizer,
          criterion,
          device,
          device_ids,
          nas_config,
          lookup_table,
          dataloader,
          val_dataloader,
          n_epochs=args.end_epoch,
          logger=logging,
      )
      trainer.train()
  
file_path: nni/examples/nas/legacy/oneshot/pfld/train.py
imports: from __future__ import absolute_import, division, print_function
  import argparse
  import logging
  import os
  import torch
  import torchvision
  import numpy
  from datasets import PFLDDatasets
  from lib.builder import search_space
  from lib.ops import PRIMITIVES
  from lib.trainer import PFLDTrainer
  from lib.utils import PFLDLoss
  from nni.algorithms.nas.pytorch.fbnet import LookUpTable, NASConfig
  from torch.utils.data import DataLoader
label_desc: Use image node_id 1 for calling a global function with example usage: main(args) without return types