code: import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from operations import FactorizedReduce, StdConv, MixedOp
  
  
  class Cell(nn.Module):
      """ Cell for search
      Each edge is mixed and continuous relaxed.
      """
  
      def __init__(self, num_nodes, c_prev_prev, c_prev, c_cur, reduction_prev, reduction_cur, search_space):
          """
          Args:
              num_nodes: Number of intermediate cell nodes
              c_prev_prev: channels_out[k-2]
              c_prev : Channels_out[k-1]
              c_cur   : Channels_in[k] (current)
              reduction_prev: flag for whether the previous cell is reduction cell or not
              reduction_cur: flag for whether the current cell is reduction cell or not
          """
  
          super(Cell, self).__init__()
          self.reduction_cur = reduction_cur
          self.num_nodes = num_nodes
  
          # If previous cell is reduction cell, current input size does not match with
          # output size of cell[k-2]. So the output[k-2] should be reduced by preprocessing
          if reduction_prev:
              self.preprocess0 = FactorizedReduce(c_prev_prev, c_cur)
          else:
              self.preprocess0 = StdConv(c_prev_prev, c_cur, kernel_size=1, stride=1, padding=0)
          self.preprocess1 = StdConv(c_prev, c_cur, kernel_size=1, stride=1, padding=0)
  
          # Generate dag from mixed operations
          self.dag_ops = nn.ModuleList()
  
          for i in range(self.num_nodes):
              self.dag_ops.append(nn.ModuleList())
              # Include 2 input nodes
              for j in range(2+i):
                  # Reduction with stride = 2 must be only for the input node
                  stride = 2 if reduction_cur and j < 2 else 1
                  op = MixedOp(c_cur, stride, search_space)
                  self.dag_ops[i].append(op)
  
      def forward(self, s0, s1, w_dag):
          s0 = self.preprocess0(s0)
          s1 = self.preprocess1(s1)
  
          states = [s0, s1]
          for edges, w_list in zip(self.dag_ops, w_dag):
              state_cur = sum(edges[i](s, w) for i, (s, w) in enumerate((zip(states, w_list))))
              states.append(state_cur)
  
          state_out = torch.cat(states[2:], dim=1)
          return state_out
  
  
  class NetworkCNN(nn.Module):
  
      def __init__(self, init_channels, input_channels, num_classes,
                   num_layers, criterion, search_space, num_nodes, stem_multiplier):
          super(NetworkCNN, self).__init__()
  
          self.init_channels = init_channels
          self.num_classes = num_classes
          self.num_layers = num_layers
          self.criterion = criterion
  
          # TODO: Algorithm settings?
          self.num_nodes = num_nodes
          self.stem_multiplier = stem_multiplier
  
          c_cur = self.stem_multiplier*self.init_channels
  
          self.stem = nn.Sequential(
              nn.Conv2d(input_channels, c_cur, 3, padding=1, bias=False),
              nn.BatchNorm2d(c_cur)
          )
  
          # In first Cell stem is used for s0 and s1
          # c_prev_prev and c_prev - output channels size
          # c_cur - init channels size
          c_prev_prev, c_prev, c_cur = c_cur, c_cur, self.init_channels
  
          self.cells = nn.ModuleList()
  
          reduction_prev = False
          for i in range(self.num_layers):
              # For Network with 1 layer: Only Normal Cell
              if self.num_layers == 1:
                  reduction_cur = False
              else:
                  # For Network with two layers: First layer - Normal, Second - Reduction
                  # For Other Networks: [1/3, 2/3] Layers - Reduction cell with double channels
                  # Others - Normal cell
                  if ((self.num_layers == 2 and i == 1) or
                          (self.num_layers > 2 and i in [self.num_layers//3, 2*self.num_layers//3])):
                      c_cur *= 2
                      reduction_cur = True
                  else:
                      reduction_cur = False
  
              cell = Cell(self.num_nodes, c_prev_prev, c_prev, c_cur, reduction_prev, reduction_cur, search_space)
              reduction_prev = reduction_cur
              self.cells.append(cell)
  
              c_cur_out = c_cur * self.num_nodes
              c_prev_prev, c_prev = c_prev, c_cur_out
  
          self.global_pooling = nn.AdaptiveAvgPool2d(1)
          self.classifier = nn.Linear(c_prev, self.num_classes)
  
          # Initialize alphas parameters
          num_ops = len(search_space.primitives)
  
          self.alpha_normal = nn.ParameterList()
          self.alpha_reduce = nn.ParameterList()
  
          for i in range(self.num_nodes):
              self.alpha_normal.append(nn.Parameter(1e-3*torch.randn(i+2, num_ops)))
              if self.num_layers > 1:
                  self.alpha_reduce.append(nn.Parameter(1e-3*torch.randn(i+2, num_ops)))
  
          # Setup alphas list
          self.alphas = []
          for name, parameter in self.named_parameters():
              if "alpha" in name:
                  self.alphas.append((name, parameter))
  
      def forward(self, x):
  
          weights_normal = [F.softmax(alpha, dim=-1) for alpha in self.alpha_normal]
          weights_reduce = [F.softmax(alpha, dim=-1) for alpha in self.alpha_reduce]
  
          s0 = s1 = self.stem(x)
  
          for cell in self.cells:
              weights = weights_reduce if cell.reduction_cur else weights_normal
              s0, s1 = s1, cell(s0, s1, weights)
  
          out = self.global_pooling(s1)
  
          # Make out flatten
          out = out.view(out.size(0), -1)
  
          logits = self.classifier(out)
          return logits
  
      def print_alphas(self):
  
          print("\n>>> Alphas Normal <<<")
          for alpha in self.alpha_normal:
              print(F.softmax(alpha, dim=-1))
  
          if self.num_layers > 1:
              print("\n>>> Alpha Reduce <<<")
              for alpha in self.alpha_reduce:
                  print(F.softmax(alpha, dim=-1))
          print("\n")
  
      def getWeights(self):
          return self.parameters()
  
      def getAlphas(self):
          for _, parameter in self.alphas:
              yield parameter
  
      def loss(self, x, y):
          logits = self.forward(x)
          return self.criterion(logits, y)
  
      def genotype(self, search_space):
          gene_normal = search_space.parse(self.alpha_normal, k=2)
          gene_reduce = search_space.parse(self.alpha_reduce, k=2)
          # concat all intermediate nodes
          concat = range(2, 2 + self.num_nodes)
  
          return search_space.genotype(normal=gene_normal, normal_concat=concat,
                                       reduce=gene_reduce, reduce_concat=concat)
  
apis: torch.nn.BatchNorm2d torch.nn.Linear torch.randn torch.nn.AdaptiveAvgPool2d torch.nn.functional.softmax torch.nn.ParameterList torch.nn.ModuleList torch.nn.Conv2d torch.cat