code: def __init__(
      self,
      dataset,
      n_layers,
      in_channels=3,
      channels=16,
      n_nodes=4,
      retrain=False,
      shared_modules=None,
  ):
      super().__init__()
      assert dataset in ["cifar10", "imagenet"]
      self.dataset = dataset
      self.input_size = 32 if dataset == "cifar" else 224
      self.in_channels = in_channels
      self.channels = channels
      self.n_nodes = n_nodes
      self.aux_size = {2 * n_layers // 3: self.input_size // 4}
      if dataset == "cifar10":
          self.n_classes = 10
          self.aux_head_class = (
              AuxiliaryHeadCIFAR if retrain else DistillHeadCIFAR
          )
          if not retrain:
              self.aux_size = {n_layers // 3: 6, 2 * n_layers // 3: 6}
      elif dataset == "imagenet":
          self.n_classes = 1000
          self.aux_head_class = (
              AuxiliaryHeadImageNet if retrain else DistillHeadImagenet
          )
          if not retrain:
              self.aux_size = {n_layers // 3: 6, 2 * n_layers // 3: 5}
      self.n_layers = n_layers
      self.aux_head = nn.ModuleDict()
      self.ensemble_param = (
          nn.Parameter(
              torch.rand(len(self.aux_size) + 1)
              / (len(self.aux_size) + 1)
          )
          if not retrain
          else None
      )
  
      stem_multiplier = 3 if dataset == "cifar" else 1
      c_cur = stem_multiplier * self.channels
      self.shared_modules = {}  # do not wrap with ModuleDict
      if shared_modules is not None:
          self.stem = shared_modules["stem"]
      else:
          self.stem = nn.Sequential(
              nn.Conv2d(in_channels, c_cur, 3, 1, 1, bias=False),
              nn.BatchNorm2d(c_cur),
          )
          self.shared_modules["stem"] = self.stem
  
      # for the first cell, stem is used for both s0 and s1
      # [!] channels_pp and channels_p is output channel size, but c_cur is input channel size.
      channels_pp, channels_p, c_cur = c_cur, c_cur, channels
  
      self.cells = nn.ModuleList()
      reduction_p, reduction = False, False
      aux_head_count = 0
      for i in range(n_layers):
          reduction_p, reduction = reduction, False
          if i in [n_layers // 3, 2 * n_layers // 3]:
              c_cur *= 2
              reduction = True
  
          cell = Cell(
              n_nodes,
              channels_pp,
              channels_p,
              c_cur,
              reduction_p,
              reduction,
          )
          self.cells.append(cell)
          c_cur_out = c_cur * n_nodes
          if i in self.aux_size:
              if shared_modules is not None:
                  self.aux_head[str(i)] = shared_modules[
                      "aux" + str(aux_head_count)
                  ]
              else:
                  self.aux_head[str(i)] = self.aux_head_class(
                      c_cur_out, self.aux_size[i], self.n_classes
                  )
                  self.shared_modules[
                      "aux" + str(aux_head_count)
                  ] = self.aux_head[str(i)]
              aux_head_count += 1
          channels_pp, channels_p = channels_p, c_cur_out
  
      self.gap = nn.AdaptiveAvgPool2d(1)
      self.linear = nn.Linear(channels_p, self.n_classes)
  
imports: import torch
  import torch.nn
  import torch.nn.functional
  import ops
  import numpy
  from nni.nas.pytorch import mutables
  from utils import parse_results
  from aux_head import DistillHeadCIFAR, DistillHeadImagenet, AuxiliaryHeadCIFAR, AuxiliaryHeadImageNet
label_desc: Use image node_id 1 to create a new Model object from inherited base classes: nn with example: obj = Model(dataset, n_layers, in_channels, channels, n_nodes, retrain, shared_modules)