code: def __init__(
      self,
      arch_str,
      num_classes,
      use_bn=True,
      keep_mask=None,
      stem_ch=16,
  ):
      super(NAS201Model, self).__init__()
      self.arch_str = arch_str
      self.num_classes = num_classes
      self.use_bn = use_bn
      self.stem_ch = stem_ch
  
      self.stem = stem(out_channels=stem_ch, use_bn=use_bn)
      self.stack_cell1 = nn.Sequential(
          *[
              SearchCell(
                  in_channels=stem_ch,
                  out_channels=stem_ch,
                  stride=1,
                  affine=False,
                  track_running_stats=False,
                  use_bn=use_bn,
                  keep_mask=keep_mask,
              )
              for i in range(5)
          ]
      )
      self.reduction1 = reduction(
          in_channels=stem_ch, out_channels=stem_ch * 2
      )
      self.stack_cell2 = nn.Sequential(
          *[
              SearchCell(
                  in_channels=stem_ch * 2,
                  out_channels=stem_ch * 2,
                  stride=1,
                  affine=False,
                  track_running_stats=False,
                  use_bn=use_bn,
                  keep_mask=keep_mask,
              )
              for i in range(5)
          ]
      )
      self.reduction2 = reduction(
          in_channels=stem_ch * 2, out_channels=stem_ch * 4
      )
      self.stack_cell3 = nn.Sequential(
          *[
              SearchCell(
                  in_channels=stem_ch * 4,
                  out_channels=stem_ch * 4,
                  stride=1,
                  affine=False,
                  track_running_stats=False,
                  use_bn=use_bn,
                  keep_mask=keep_mask,
              )
              for i in range(5)
          ]
      )
      self.top = top(
          in_dims=stem_ch * 4, num_classes=num_classes, use_bn=use_bn
      )
  
file_path: singa/examples/model_selection/TRAILS-Database-Native-Model-Selection/internal/ml/model_selection/src/third_pkg/sp201_lib/nasbench2.py
imports: from .nasbench2_ops import *
label_desc: Use image node_id 1 to create a new NAS201Model object from inherited base classes: nn with example: obj = NAS201Model(arch_str, num_classes, use_bn, keep_mask, stem_ch)