update GDAS and SETN

lib/models/cell_searchs/__init__.py

@@ -0,0 +1,4 @@
+from .search_model_darts_v1 import TinyNetworkDartsV1
+from .search_model_darts_v2 import TinyNetworkDartsV2
+from .search_model_gdas     import TinyNetworkGDAS
+from .search_model_setn     import TinyNetworkSETN

lib/models/cell_searchs/cells.py

@@ -2,7 +2,7 @@ import math, torch
 import torch.nn as nn
 import torch.nn.functional as F
 from copy import deepcopy
-from .operations import OPS, ReLUConvBN
+from ..cell_operations import OPS
 
 
 class SearchCell(nn.Module):
@@ -113,84 +113,3 @@ class SearchCell(nn.Module):
         inter_nodes.append( self.edges[node_str][op_index]( nodes[j] ) )
       nodes.append( sum(inter_nodes) )
     return nodes[-1]
-
-
-class InferCell(nn.Module):
-
-  def __init__(self, genotype, C_in, C_out, stride):
-    super(InferCell, self).__init__()
-
-    self.layers  = nn.ModuleList()
-    self.node_IN = []
-    self.node_IX = []
-    self.genotype = deepcopy(genotype)
-    for i in range(1, len(genotype)):
-      node_info = genotype[i-1]
-      cur_index = []
-      cur_innod = []
-      for (op_name, op_in) in node_info:
-        if op_in == 0:
-          layer = OPS[op_name](C_in , C_out, stride)
-        else:
-          layer = OPS[op_name](C_out, C_out,      1)
-        cur_index.append( len(self.layers) )
-        cur_innod.append( op_in )
-        self.layers.append( layer )
-      self.node_IX.append( cur_index )
-      self.node_IN.append( cur_innod )
-    self.nodes   = len(genotype)
-    self.in_dim  = C_in
-    self.out_dim = C_out
-
-  def extra_repr(self):
-    string = 'info :: nodes={nodes}, inC={in_dim}, outC={out_dim}'.format(**self.__dict__)
-    laystr = []
-    for i, (node_layers, node_innods) in enumerate(zip(self.node_IX,self.node_IN)):
-      y = ['I{:}-L{:}'.format(_ii, _il) for _il, _ii in zip(node_layers, node_innods)]
-      x = '{:}<-({:})'.format(i+1, ','.join(y))
-      laystr.append( x )
-    return string + ', [{:}]'.format( ' | '.join(laystr) ) + ', {:}'.format(self.genotype.tostr())
-
-  def forward(self, inputs):
-    nodes = [inputs]
-    for i, (node_layers, node_innods) in enumerate(zip(self.node_IX,self.node_IN)):
-      node_feature = sum( self.layers[_il](nodes[_ii]) for _il, _ii in zip(node_layers, node_innods) )
-      nodes.append( node_feature )
-    return nodes[-1]
-
-
-
-class ResNetBasicblock(nn.Module):
-
-  def __init__(self, inplanes, planes, stride):
-    super(ResNetBasicblock, self).__init__()
-    assert stride == 1 or stride == 2, 'invalid stride {:}'.format(stride)
-    self.conv_a = ReLUConvBN(inplanes, planes, 3, stride, 1, 1)
-    self.conv_b = ReLUConvBN(  planes, planes, 3,      1, 1, 1)
-    if stride == 2:
-      self.downsample = nn.Sequential(
-                           nn.AvgPool2d(kernel_size=2, stride=2, padding=0),
-                           nn.Conv2d(inplanes, planes, kernel_size=1, stride=1, padding=0, bias=False))
-    elif inplanes != planes:
-      self.downsample = ReLUConvBN(inplanes, planes, 1, 1, 0, 1)
-    else:
-      self.downsample = None
-    self.in_dim  = inplanes
-    self.out_dim = planes
-    self.stride  = stride
-    self.num_conv = 2
-
-  def extra_repr(self):
-    string = '{name}(inC={in_dim}, outC={out_dim}, stride={stride})'.format(name=self.__class__.__name__, **self.__dict__)
-    return string
-
-  def forward(self, inputs):
-
-    basicblock = self.conv_a(inputs)
-    basicblock = self.conv_b(basicblock)
-
-    if self.downsample is not None:
-      residual = self.downsample(inputs)
-    else:
-      residual = inputs
-    return residual + basicblock

lib/models/cell_searchs/genotypes.py

@@ -0,0 +1,158 @@
+from copy import deepcopy
+
+
+
+def get_combination(space, num):
+  combs = []
+  for i in range(num):
+    if i == 0:
+      for func in space:
+        combs.append( [(func, i)] )
+    else:
+      new_combs = []
+      for string in combs:
+        for func in space:
+          xstring = string + [(func, i)]
+          new_combs.append( xstring )
+      combs = new_combs
+  return combs
+  
+
+
+class Structure:
+
+  def __init__(self, genotype):
+    assert isinstance(genotype, list) or isinstance(genotype, tuple), 'invalid class of genotype : {:}'.format(type(genotype))
+    self.node_num = len(genotype) + 1
+    self.nodes    = []
+    self.node_N   = []
+    for idx, node_info in enumerate(genotype):
+      assert isinstance(node_info, list) or isinstance(node_info, tuple), 'invalid class of node_info : {:}'.format(type(node_info))
+      assert len(node_info) >= 1, 'invalid length : {:}'.format(len(node_info))
+      for node_in in node_info:
+        assert isinstance(node_in, list) or isinstance(node_in, tuple), 'invalid class of in-node : {:}'.format(type(node_in))
+        assert len(node_in) == 2 and node_in[1] <= idx, 'invalid in-node : {:}'.format(node_in)
+      self.node_N.append( len(node_info) )
+      self.nodes.append( tuple(deepcopy(node_info)) )
+
+  def tolist(self, remove_str):
+    # convert this class to the list, if remove_str is 'none', then remove the 'none' operation.
+    # note that we re-order the input node in this function
+    # return the-genotype-list and success [if unsuccess, it is not a connectivity]
+    genotypes = []
+    for node_info in self.nodes:
+      node_info = list( node_info )
+      node_info = sorted(node_info, key=lambda x: (x[1], x[0]))
+      node_info = tuple(filter(lambda x: x[0] != remove_str, node_info))
+      if len(node_info) == 0: return None, False
+      genotypes.append( node_info )
+    return genotypes, True
+
+  def node(self, index):
+    assert index > 0 and index <= len(self), 'invalid index={:} < {:}'.format(index, len(self))
+    return self.nodes[index]
+
+  def tostr(self):
+    strings = []
+    for node_info in self.nodes:
+      string = '|'.join([x[0]+'~{:}'.format(x[1]) for x in node_info])
+      string = '|{:}|'.format(string)
+      strings.append( string )
+    return '+'.join(strings)
+
+  def __repr__(self):
+    return ('{name}({node_num} nodes with {node_info})'.format(name=self.__class__.__name__, node_info=self.tostr(), **self.__dict__))
+
+  def __len__(self):
+    return len(self.nodes) + 1
+
+  def __getitem__(self, index):
+    return self.nodes[index]
+
+  @staticmethod
+  def str2structure(xstr):
+    assert isinstance(xstr, str), 'must take string (not {:}) as input'.format(type(xstr))
+    nodestrs = xstr.split('+')
+    genotypes = []
+    for i, node_str in enumerate(nodestrs):
+      inputs = list(filter(lambda x: x != '', node_str.split('|')))
+      for xinput in inputs: assert len(xinput.split('~')) == 2, 'invalid input length : {:}'.format(xinput)
+      inputs = ( xi.split('~') for xi in inputs )
+      input_infos = tuple( (op, int(IDX)) for (op, IDX) in inputs)
+      genotypes.append( input_infos )
+    return Structure( genotypes )
+
+  @staticmethod
+  def str2fullstructure(xstr, default_name='none'):
+    assert isinstance(xstr, str), 'must take string (not {:}) as input'.format(type(xstr))
+    nodestrs = xstr.split('+')
+    genotypes = []
+    for i, node_str in enumerate(nodestrs):
+      inputs = list(filter(lambda x: x != '', node_str.split('|')))
+      for xinput in inputs: assert len(xinput.split('~')) == 2, 'invalid input length : {:}'.format(xinput)
+      inputs = ( xi.split('~') for xi in inputs )
+      input_infos = list( (op, int(IDX)) for (op, IDX) in inputs)
+      all_in_nodes= list(x[1] for x in input_infos)
+      for j in range(i):
+        if j not in all_in_nodes: input_infos.append((default_name, j))
+      node_info = sorted(input_infos, key=lambda x: (x[1], x[0]))
+      genotypes.append( tuple(node_info) )
+    return Structure( genotypes )
+
+  @staticmethod
+  def gen_all(search_space, num, return_ori):
+    assert isinstance(search_space, list) or isinstance(search_space, tuple), 'invalid class of search-space : {:}'.format(type(search_space))
+    assert num >= 2, 'There should be at least two nodes in a neural cell instead of {:}'.format(num)
+    all_archs = get_combination(search_space, 1)
+    for i, arch in enumerate(all_archs):
+      all_archs[i] = [ tuple(arch) ]
+  
+    for inode in range(2, num):
+      cur_nodes = get_combination(search_space, inode)
+      new_all_archs = []
+      for previous_arch in all_archs:
+        for cur_node in cur_nodes:
+          new_all_archs.append( previous_arch + [tuple(cur_node)] )
+      all_archs = new_all_archs
+    if return_ori:
+      return all_archs
+    else:
+      return [Structure(x) for x in all_archs]
+
+
+
+ResNet_CODE = Structure(
+  [(('nor_conv_3x3', 0), ), # node-1 
+   (('nor_conv_3x3', 1), ), # node-2
+   (('skip_connect', 0), ('skip_connect', 2))] # node-3
+  )
+
+AllConv3x3_CODE = Structure(
+  [(('nor_conv_3x3', 0), ), # node-1 
+   (('nor_conv_3x3', 0), ('nor_conv_3x3', 1)), # node-2
+   (('nor_conv_3x3', 0), ('nor_conv_3x3', 1), ('nor_conv_3x3', 2))] # node-3
+  )
+
+AllFull_CODE = Structure(
+  [(('skip_connect', 0), ('nor_conv_1x1', 0), ('nor_conv_3x3', 0), ('avg_pool_3x3', 0)), # node-1 
+   (('skip_connect', 0), ('nor_conv_1x1', 0), ('nor_conv_3x3', 0), ('avg_pool_3x3', 0), ('skip_connect', 1), ('nor_conv_1x1', 1), ('nor_conv_3x3', 1), ('avg_pool_3x3', 1)), # node-2
+   (('skip_connect', 0), ('nor_conv_1x1', 0), ('nor_conv_3x3', 0), ('avg_pool_3x3', 0), ('skip_connect', 1), ('nor_conv_1x1', 1), ('nor_conv_3x3', 1), ('avg_pool_3x3', 1), ('skip_connect', 2), ('nor_conv_1x1', 2), ('nor_conv_3x3', 2), ('avg_pool_3x3', 2))] # node-3
+  )
+
+AllConv1x1_CODE = Structure(
+  [(('nor_conv_1x1', 0), ), # node-1 
+   (('nor_conv_1x1', 0), ('nor_conv_1x1', 1)), # node-2
+   (('nor_conv_1x1', 0), ('nor_conv_1x1', 1), ('nor_conv_1x1', 2))] # node-3
+  )
+
+AllIdentity_CODE = Structure(
+  [(('skip_connect', 0), ), # node-1 
+   (('skip_connect', 0), ('skip_connect', 1)), # node-2
+   (('skip_connect', 0), ('skip_connect', 1), ('skip_connect', 2))] # node-3
+  )
+
+architectures = {'resnet'  : ResNet_CODE,
+                 'all_c3x3': AllConv3x3_CODE,
+                 'all_c1x1': AllConv1x1_CODE,
+                 'all_idnt': AllIdentity_CODE,
+                 'all_full': AllFull_CODE}

lib/models/cell_searchs/operations.py

@@ -1,162 +0,0 @@
-import torch
-import torch.nn as nn
-
-__all__ = ['OPS', 'ReLUConvBN', 'SearchSpaceNames']
-
-OPS = {
-  'none'         : lambda C_in, C_out, stride: Zero(C_in, C_out, stride),
-  'avg_pool_3x3' : lambda C_in, C_out, stride: POOLING(C_in, C_out, stride, 'avg'),
-  'max_pool_3x3' : lambda C_in, C_out, stride: POOLING(C_in, C_out, stride, 'max'),
-  'nor_conv_7x7' : lambda C_in, C_out, stride: ReLUConvBN(C_in, C_out, (7,7), (stride,stride), (3,3), (1,1)),
-  'nor_conv_3x3' : lambda C_in, C_out, stride: ReLUConvBN(C_in, C_out, (3,3), (stride,stride), (1,1), (1,1)),
-  'nor_conv_1x1' : lambda C_in, C_out, stride: ReLUConvBN(C_in, C_out, (1,1), (stride,stride), (0,0), (1,1)),
-  'skip_connect' : lambda C_in, C_out, stride: Identity() if stride == 1 and C_in == C_out else FactorizedReduce(C_in, C_out, stride),
-}
-
-CONNECT_NAS_BENCHMARK  = ['none', 'skip_connect', 'nor_conv_3x3']
-
-SearchSpaceNames = {'connect-nas' : CONNECT_NAS_BENCHMARK}
-
-
-class POOLING(nn.Module):
-
-  def __init__(self, C_in, C_out, stride, mode):
-    super(POOLING, self).__init__()
-    if C_in == C_out:
-      self.preprocess = None
-    else:
-      self.preprocess = ReLUConvBN(C_in, C_out, 1, 1, 0)
-    if mode == 'avg'  : self.op = nn.AvgPool2d(3, stride=stride, padding=1, count_include_pad=False)
-    elif mode == 'max': self.op = nn.MaxPool2d(3, stride=stride, padding=1)
-    else              : raise ValueError('Invalid mode={:} in POOLING'.format(mode))
-
-  def forward(self, inputs):
-    if self.preprocess: x = self.preprocess(inputs)
-    else              : x = inputs
-    return self.op(x)
-
-
-class ReLUConvBN(nn.Module):
-
-  def __init__(self, C_in, C_out, kernel_size, stride, padding, dilation):
-    super(ReLUConvBN, self).__init__()
-    self.op = nn.Sequential(
-      nn.ReLU(inplace=False),
-      nn.Conv2d(C_in, C_out, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=False),
-      nn.BatchNorm2d(C_out)
-    )
-
-  def forward(self, x):
-    return self.op(x)
-
-
-class Identity(nn.Module):
-
-  def __init__(self):
-    super(Identity, self).__init__()
-
-  def forward(self, x):
-    return x
-
-
-class Zero(nn.Module):
-
-  def __init__(self, C_in, C_out, stride):
-    super(Zero, self).__init__()
-    self.C_in   = C_in
-    self.C_out  = C_out
-    self.stride = stride
-    self.is_zero = True
-
-  def forward(self, x):
-    if self.C_in == self.C_out:
-      if self.stride == 1: return x.mul(0.)
-      else               : return x[:,:,::self.stride,::self.stride].mul(0.)
-    else:
-      shape = list(x.shape)
-      shape[1] = self.C_out
-      zeros = x.new_zeros(shape, dtype=x.dtype, device=x.device)
-      return zeros
-
-  def extra_repr(self):
-    return 'C_in={C_in}, C_out={C_out}, stride={stride}'.format(**self.__dict__)
-
-
-class FactorizedReduce(nn.Module):
-
-  def __init__(self, C_in, C_out, stride):
-    super(FactorizedReduce, self).__init__()
-    self.stride = stride
-    self.C_in   = C_in  
-    self.C_out  = C_out  
-    self.relu   = nn.ReLU(inplace=False)
-    if stride == 2:
-      #assert C_out % 2 == 0, 'C_out : {:}'.format(C_out)
-      C_outs = [C_out // 2, C_out - C_out // 2]
-      self.convs = nn.ModuleList()
-      for i in range(2):
-        self.convs.append( nn.Conv2d(C_in, C_outs[i], 1, stride=stride, padding=0, bias=False) )
-      self.pad = nn.ConstantPad2d((0, 1, 0, 1), 0)
-    else:
-      raise ValueError('Invalid stride : {:}'.format(stride))
-    
-    self.bn = nn.BatchNorm2d(C_out)
-
-  def forward(self, x):
-    x = self.relu(x)
-    y = self.pad(x)
-    out = torch.cat([self.convs[0](x), self.convs[1](y[:,:,1:,1:])], dim=1)
-    out = self.bn(out)
-    return out
-
-  def extra_repr(self):
-    return 'C_in={C_in}, C_out={C_out}, stride={stride}'.format(**self.__dict__)
-
-
-# Auto-ReID: Searching for a Part-Aware ConvNet for Person Re-Identification, ICCV 2019
-class PartAwareOp(nn.Module):
-
-  def __init__(self, C_in, C_out, stride, part=4):
-    super().__init__()
-    self.part   = 4
-    self.hidden = C_in // 3
-    self.avg_pool = nn.AdaptiveAvgPool2d(1)
-    self.local_conv_list = nn.ModuleList()
-    for i in range(self.part):
-      self.local_conv_list.append(
-            nn.Sequential(nn.ReLU(), nn.Conv2d(C_in, self.hidden, 1), nn.BatchNorm2d(self.hidden, affine=True))
-            )
-    self.W_K = nn.Linear(self.hidden, self.hidden)
-    self.W_Q = nn.Linear(self.hidden, self.hidden)
-
-    if stride == 2  : self.last = FactorizedReduce(C_in + self.hidden, C_out, 2)
-    elif stride == 1: self.last = FactorizedReduce(C_in + self.hidden, C_out, 1)
-    else:             raise ValueError('Invalid Stride : {:}'.format(stride))
-
-  def forward(self, x):
-    batch, C, H, W = x.size()
-    assert H >= self.part, 'input size too small : {:} vs {:}'.format(x.shape, self.part)
-    IHs = [0]
-    for i in range(self.part): IHs.append( min(H, int((i+1)*(float(H)/self.part))) )
-    local_feat_list = []
-    for i in range(self.part):
-      feature = x[:, :, IHs[i]:IHs[i+1], :]
-      xfeax   = self.avg_pool(feature)
-      xfea    = self.local_conv_list[i]( xfeax )
-      local_feat_list.append( xfea )
-    part_feature = torch.cat(local_feat_list, dim=2).view(batch, -1, self.part)
-    part_feature = part_feature.transpose(1,2).contiguous()
-    part_K       = self.W_K(part_feature)
-    part_Q       = self.W_Q(part_feature).transpose(1,2).contiguous()
-    weight_att   = torch.bmm(part_K, part_Q)
-    attention    = torch.softmax(weight_att, dim=2)
-    aggreateF    = torch.bmm(attention, part_feature).transpose(1,2).contiguous()
-    features = []
-    for i in range(self.part):
-      feature = aggreateF[:, :, i:i+1].expand(batch, self.hidden, IHs[i+1]-IHs[i])
-      feature = feature.view(batch, self.hidden, IHs[i+1]-IHs[i], 1)
-      features.append( feature )
-    features  = torch.cat(features, dim=2).expand(batch, self.hidden, H, W)
-    final_fea = torch.cat((x,features), dim=1)
-    outputs   = self.last( final_fea )
-    return outputs

lib/models/cell_searchs/search_cells.py

@@ -0,0 +1,134 @@
+import math, random, torch
+import warnings
+import torch.nn as nn
+import torch.nn.functional as F
+from copy import deepcopy
+from ..cell_operations import OPS
+
+
+class SearchCell(nn.Module):
+
+  def __init__(self, C_in, C_out, stride, max_nodes, op_names):
+    super(SearchCell, self).__init__()
+
+    self.op_names  = deepcopy(op_names)
+    self.edges     = nn.ModuleDict()
+    self.max_nodes = max_nodes
+    self.in_dim    = C_in
+    self.out_dim   = C_out
+    for i in range(1, max_nodes):
+      for j in range(i):
+        node_str = '{:}<-{:}'.format(i, j)
+        if j == 0:
+          xlists = [OPS[op_name](C_in , C_out, stride) for op_name in op_names]
+        else:
+          xlists = [OPS[op_name](C_in , C_out,      1) for op_name in op_names]
+        self.edges[ node_str ] = nn.ModuleList( xlists )
+    self.edge_keys  = sorted(list(self.edges.keys()))
+    self.edge2index = {key:i for i, key in enumerate(self.edge_keys)}
+    self.num_edges  = len(self.edges)
+
+  def extra_repr(self):
+    string = 'info :: {max_nodes} nodes, inC={in_dim}, outC={out_dim}'.format(**self.__dict__)
+    return string
+
+  def forward(self, inputs, weightss):
+    nodes = [inputs]
+    for i in range(1, self.max_nodes):
+      inter_nodes = []
+      for j in range(i):
+        node_str = '{:}<-{:}'.format(i, j)
+        weights  = weightss[ self.edge2index[node_str] ]
+        inter_nodes.append( sum( layer(nodes[j]) * w for layer, w in zip(self.edges[node_str], weights) ) )
+      nodes.append( sum(inter_nodes) )
+    return nodes[-1]
+
+  # GDAS
+  def forward_gdas(self, inputs, alphas, _tau):
+    avoid_zero = 0
+    while True:
+      gumbels = -torch.empty_like(alphas).exponential_().log()
+      logits  = (alphas.log_softmax(dim=1) + gumbels) / _tau
+      probs   = nn.functional.softmax(logits, dim=1)
+      index   = probs.max(-1, keepdim=True)[1]
+      one_h   = torch.zeros_like(logits).scatter_(-1, index, 1.0)
+      hardwts = one_h - probs.detach() + probs
+      if (torch.isinf(gumbels).any()) or (torch.isinf(probs).any()) or (torch.isnan(probs).any()):
+        continue # avoid the numerical error
+      nodes   = [inputs]
+      for i in range(1, self.max_nodes):
+        inter_nodes = []
+        for j in range(i):
+          node_str = '{:}<-{:}'.format(i, j)
+          weights  = hardwts[ self.edge2index[node_str] ]
+          argmaxs  = index[ self.edge2index[node_str] ].item()
+          weigsum  = sum( weights[_ie] * edge(nodes[j]) if _ie == argmaxs else weights[_ie] for _ie, edge in enumerate(self.edges[node_str]) )
+          inter_nodes.append( weigsum )
+        nodes.append( sum(inter_nodes) )
+      avoid_zero += 1
+      if nodes[-1].sum().item() == 0:
+        if avoid_zero < 10: continue
+        else:
+          warnings.warn('get zero outputs with avoid_zero={:}'.format(avoid_zero))
+          break
+      else:
+        break
+    return nodes[-1]
+
+  # joint
+  def forward_joint(self, inputs, weightss):
+    nodes = [inputs]
+    for i in range(1, self.max_nodes):
+      inter_nodes = []
+      for j in range(i):
+        node_str = '{:}<-{:}'.format(i, j)
+        weights  = weightss[ self.edge2index[node_str] ]
+        aggregation = sum( layer(nodes[j]) * w for layer, w in zip(self.edges[node_str], weights) ) / weights.numel()
+        inter_nodes.append( aggregation )
+      nodes.append( sum(inter_nodes) )
+    return nodes[-1]
+
+  # uniform random sampling per iteration
+  def forward_urs(self, inputs):
+    nodes = [inputs]
+    for i in range(1, self.max_nodes):
+      while True: # to avoid select zero for all ops
+        sops, has_non_zero = [], False
+        for j in range(i):
+          node_str   = '{:}<-{:}'.format(i, j)
+          candidates = self.edges[node_str]
+          select_op  = random.choice(candidates)
+          sops.append( select_op )
+          if not hasattr(select_op, 'is_zero') or select_op.is_zero == False: has_non_zero=True
+        if has_non_zero: break
+      inter_nodes = []
+      for j, select_op in enumerate(sops):
+        inter_nodes.append( select_op(nodes[j]) )
+      nodes.append( sum(inter_nodes) )
+    return nodes[-1]
+
+  # select the argmax
+  def forward_select(self, inputs, weightss):
+    nodes = [inputs]
+    for i in range(1, self.max_nodes):
+      inter_nodes = []
+      for j in range(i):
+        node_str = '{:}<-{:}'.format(i, j)
+        weights  = weightss[ self.edge2index[node_str] ]
+        inter_nodes.append( self.edges[node_str][ weights.argmax().item() ]( nodes[j] ) )
+        #inter_nodes.append( sum( layer(nodes[j]) * w for layer, w in zip(self.edges[node_str], weights) ) )
+      nodes.append( sum(inter_nodes) )
+    return nodes[-1]
+
+  # forward with a specific structure
+  def forward_dynamic(self, inputs, structure):
+    nodes = [inputs]
+    for i in range(1, self.max_nodes):
+      cur_op_node = structure.nodes[i-1]
+      inter_nodes = []
+      for op_name, j in cur_op_node:
+        node_str = '{:}<-{:}'.format(i, j)
+        op_index = self.op_names.index( op_name )
+        inter_nodes.append( self.edges[node_str][op_index]( nodes[j] ) )
+      nodes.append( sum(inter_nodes) )
+    return nodes[-1]

lib/models/cell_searchs/search_model_darts_v1.py

@@ -0,0 +1,93 @@
+##################################################
+# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2019 #
+########################################################
+# DARTS: Differentiable Architecture Search, ICLR 2019 #
+########################################################
+import torch
+import torch.nn as nn
+from copy import deepcopy
+from ..cell_operations import ResNetBasicblock
+from .search_cells     import SearchCell
+from .genotypes        import Structure
+
+
+class TinyNetworkDartsV1(nn.Module):
+
+  def __init__(self, C, N, max_nodes, num_classes, search_space):
+    super(TinyNetworkDartsV1, self).__init__()
+    self._C        = C
+    self._layerN   = N
+    self.max_nodes = max_nodes
+    self.stem = nn.Sequential(
+                    nn.Conv2d(3, C, kernel_size=3, padding=1, bias=False),
+                    nn.BatchNorm2d(C))
+  
+    layer_channels   = [C    ] * N + [C*2 ] + [C*2  ] * N + [C*4 ] + [C*4  ] * N    
+    layer_reductions = [False] * N + [True] + [False] * N + [True] + [False] * N
+
+    C_prev, num_edge, edge2index = C, None, None
+    self.cells = nn.ModuleList()
+    for index, (C_curr, reduction) in enumerate(zip(layer_channels, layer_reductions)):
+      if reduction:
+        cell = ResNetBasicblock(C_prev, C_curr, 2)
+      else:
+        cell = SearchCell(C_prev, C_curr, 1, max_nodes, search_space)
+        if num_edge is None: num_edge, edge2index = cell.num_edges, cell.edge2index
+        else: assert num_edge == cell.num_edges and edge2index == cell.edge2index, 'invalid {:} vs. {:}.'.format(num_edge, cell.num_edges)
+      self.cells.append( cell )
+      C_prev = cell.out_dim
+    self.op_names   = deepcopy( search_space )
+    self._Layer     = len(self.cells)
+    self.edge2index = edge2index
+    self.lastact    = nn.Sequential(nn.BatchNorm2d(C_prev), nn.ReLU(inplace=True))
+    self.global_pooling = nn.AdaptiveAvgPool2d(1)
+    self.classifier = nn.Linear(C_prev, num_classes)
+    self.arch_parameters = nn.Parameter( 1e-3*torch.randn(num_edge, len(search_space)) )
+
+  def get_weights(self):
+    xlist = list( self.stem.parameters() ) + list( self.cells.parameters() )
+    xlist+= list( self.lastact.parameters() ) + list( self.global_pooling.parameters() )
+    xlist+= list( self.classifier.parameters() )
+    return xlist
+
+  def get_alphas(self):
+    return [self.arch_parameters]
+
+  def get_message(self):
+    string = self.extra_repr()
+    for i, cell in enumerate(self.cells):
+      string += '\n {:02d}/{:02d} :: {:}'.format(i, len(self.cells), cell.extra_repr())
+    return string
+
+  def extra_repr(self):
+    return ('{name}(C={_C}, Max-Nodes={max_nodes}, N={_layerN}, L={_Layer})'.format(name=self.__class__.__name__, **self.__dict__))
+
+  def genotype(self):
+    genotypes = []
+    for i in range(1, self.max_nodes):
+      xlist = []
+      for j in range(i):
+        node_str = '{:}<-{:}'.format(i, j)
+        with torch.no_grad():
+          weights = self.arch_parameters[ self.edge2index[node_str] ]
+          op_name = self.op_names[ weights.argmax().item() ]
+        xlist.append((op_name, j))
+      genotypes.append( tuple(xlist) )
+    return Structure( genotypes )
+
+  def forward(self, inputs):
+    alphas  = nn.functional.softmax(self.arch_parameters, dim=-1)
+
+    feature = self.stem(inputs)
+    for i, cell in enumerate(self.cells):
+      if isinstance(cell, SearchCell):
+        feature = cell(feature, alphas)
+      else:
+        feature = cell(feature)
+
+    out = self.lastact(feature)
+    out = self.global_pooling( out )
+    out = out.view(out.size(0), -1)
+    logits = self.classifier(out)
+
+    return out, logits

lib/models/cell_searchs/search_model_darts_v2.py

@@ -0,0 +1,93 @@
+##################################################
+# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2019 #
+########################################################
+# DARTS: Differentiable Architecture Search, ICLR 2019 #
+########################################################
+import torch
+import torch.nn as nn
+from copy import deepcopy
+from ..cell_operations import ResNetBasicblock
+from .search_cells     import SearchCell
+from .genotypes        import Structure
+
+
+class TinyNetworkDartsV2(nn.Module):
+
+  def __init__(self, C, N, max_nodes, num_classes, search_space):
+    super(TinyNetworkDartsV2, self).__init__()
+    self._C        = C
+    self._layerN   = N
+    self.max_nodes = max_nodes
+    self.stem = nn.Sequential(
+                    nn.Conv2d(3, C, kernel_size=3, padding=1, bias=False),
+                    nn.BatchNorm2d(C))
+  
+    layer_channels   = [C    ] * N + [C*2 ] + [C*2  ] * N + [C*4 ] + [C*4  ] * N    
+    layer_reductions = [False] * N + [True] + [False] * N + [True] + [False] * N
+
+    C_prev, num_edge, edge2index = C, None, None
+    self.cells = nn.ModuleList()
+    for index, (C_curr, reduction) in enumerate(zip(layer_channels, layer_reductions)):
+      if reduction:
+        cell = ResNetBasicblock(C_prev, C_curr, 2)
+      else:
+        cell = SearchCell(C_prev, C_curr, 1, max_nodes, search_space)
+        if num_edge is None: num_edge, edge2index = cell.num_edges, cell.edge2index
+        else: assert num_edge == cell.num_edges and edge2index == cell.edge2index, 'invalid {:} vs. {:}.'.format(num_edge, cell.num_edges)
+      self.cells.append( cell )
+      C_prev = cell.out_dim
+    self.op_names   = deepcopy( search_space )
+    self._Layer     = len(self.cells)
+    self.edge2index = edge2index
+    self.lastact    = nn.Sequential(nn.BatchNorm2d(C_prev), nn.ReLU(inplace=True))
+    self.global_pooling = nn.AdaptiveAvgPool2d(1)
+    self.classifier = nn.Linear(C_prev, num_classes)
+    self.arch_parameters = nn.Parameter( 1e-3*torch.randn(num_edge, len(search_space)) )
+
+  def get_weights(self):
+    xlist = list( self.stem.parameters() ) + list( self.cells.parameters() )
+    xlist+= list( self.lastact.parameters() ) + list( self.global_pooling.parameters() )
+    xlist+= list( self.classifier.parameters() )
+    return xlist
+
+  def get_alphas(self):
+    return [self.arch_parameters]
+
+  def get_message(self):
+    string = self.extra_repr()
+    for i, cell in enumerate(self.cells):
+      string += '\n {:02d}/{:02d} :: {:}'.format(i, len(self.cells), cell.extra_repr())
+    return string
+
+  def extra_repr(self):
+    return ('{name}(C={_C}, Max-Nodes={max_nodes}, N={_layerN}, L={_Layer})'.format(name=self.__class__.__name__, **self.__dict__))
+
+  def genotype(self):
+    genotypes = []
+    for i in range(1, self.max_nodes):
+      xlist = []
+      for j in range(i):
+        node_str = '{:}<-{:}'.format(i, j)
+        with torch.no_grad():
+          weights = self.arch_parameters[ self.edge2index[node_str] ]
+          op_name = self.op_names[ weights.argmax().item() ]
+        xlist.append((op_name, j))
+      genotypes.append( tuple(xlist) )
+    return Structure( genotypes )
+
+  def forward(self, inputs):
+    alphas  = nn.functional.softmax(self.arch_parameters, dim=-1)
+
+    feature = self.stem(inputs)
+    for i, cell in enumerate(self.cells):
+      if isinstance(cell, SearchCell):
+        feature = cell(feature, alphas)
+      else:
+        feature = cell(feature)
+
+    out = self.lastact(feature)
+    out = self.global_pooling( out )
+    out = out.view(out.size(0), -1)
+    logits = self.classifier(out)
+
+    return out, logits

lib/models/cell_searchs/search_model_gdas.py

@@ -6,9 +6,9 @@
 import torch
 import torch.nn as nn
 from copy import deepcopy
-from .infer_cells  import ResNetBasicblock
-from .search_cells import SearchCell
-from .genotypes    import Structure
+from ..cell_operations import ResNetBasicblock
+from .search_cells     import SearchCell
+from .genotypes        import Structure
 
 
 class TinyNetworkGDAS(nn.Module):
@@ -44,7 +44,6 @@ class TinyNetworkGDAS(nn.Module):
     self.classifier = nn.Linear(C_prev, num_classes)
     self.arch_parameters = nn.Parameter( 1e-3*torch.randn(num_edge, len(search_space)) )
     self.tau        = 10
-    self.nan_count  = 0
 
   def get_weights(self):
     xlist = list( self.stem.parameters() ) + list( self.cells.parameters() )
@@ -52,9 +51,8 @@ class TinyNetworkGDAS(nn.Module):
     xlist+= list( self.classifier.parameters() )
     return xlist
 
-  def set_tau(self, tau, _nan_count=0):
+  def set_tau(self, tau):
     self.tau = tau
-    self.nan_count = _nan_count
 
   def get_tau(self):
     return self.tau
@@ -85,27 +83,10 @@ class TinyNetworkGDAS(nn.Module):
     return Structure( genotypes )
 
   def forward(self, inputs):
-    def gumbel_softmax(_logits, _tau):
-      while True: # a trick to avoid the gumbels bug
-        gumbels    = -torch.empty_like(_logits).exponential_().log()
-        new_logits = (_logits.log_softmax(dim=1) + gumbels) / _tau
-        probs      = nn.functional.softmax(new_logits, dim=1)
-        index      = probs.max(-1, keepdim=True)[1]
-        if index[0].item() == self.op_names.index('none') and index[3].item() == self.op_names.index('none') and index[5].item() == self.op_names.index('none'): continue
-        if index[1].item() == self.op_names.index('none') and index[2].item() == self.op_names.index('none') and index[3].item() == self.op_names.index('none') and index[4].item() == self.op_names.index('none'): continue
-        if index[3].item() == self.op_names.index('none') and index[4].item() == self.op_names.index('none') and index[5].item() == self.op_names.index('none'): continue
-        if index[3].item() == self.op_names.index('none') and index[0].item() == self.op_names.index('none') and index[1].item() == self.op_names.index('none'): continue
-        one_h      = torch.zeros_like(_logits).scatter_(-1, index, 1.0)
-        xres       = one_h - probs.detach() + probs
-        if (not torch.isinf(gumbels).any()) and (not torch.isinf(probs).any()) and (not torch.isnan(probs).any()): break
-        self.nan_count += 1
-      return xres, index
-
     feature = self.stem(inputs)
     for i, cell in enumerate(self.cells):
       if isinstance(cell, SearchCell):
-        alphas, IDX  = gumbel_softmax(self.arch_parameters, self.tau)
-        feature = cell.forward_gdas(feature, alphas, IDX.cpu())
+        feature = cell.forward_gdas(feature, self.arch_parameters, self.tau)
       else:
         feature = cell(feature)
 

lib/models/cell_searchs/search_model_setn.py

@@ -0,0 +1,130 @@
+##################################################
+# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2019 #
+######################################################################################
+# One-Shot Neural Architecture Search via Self-Evaluated Template Network, ICCV 2019 #
+######################################################################################
+import torch
+import torch.nn as nn
+from copy import deepcopy
+from ..cell_operations import ResNetBasicblock
+from .search_cells     import SearchCell
+from .genotypes        import Structure
+
+
+class TinyNetworkSETN(nn.Module):
+
+  def __init__(self, C, N, max_nodes, num_classes, search_space):
+    super(TinyNetworkSETN, self).__init__()
+    self._C        = C
+    self._layerN   = N
+    self.max_nodes = max_nodes
+    self.stem = nn.Sequential(
+                    nn.Conv2d(3, C, kernel_size=3, padding=1, bias=False),
+                    nn.BatchNorm2d(C))
+  
+    layer_channels   = [C    ] * N + [C*2 ] + [C*2  ] * N + [C*4 ] + [C*4  ] * N    
+    layer_reductions = [False] * N + [True] + [False] * N + [True] + [False] * N
+
+    C_prev, num_edge, edge2index = C, None, None
+    self.cells = nn.ModuleList()
+    for index, (C_curr, reduction) in enumerate(zip(layer_channels, layer_reductions)):
+      if reduction:
+        cell = ResNetBasicblock(C_prev, C_curr, 2)
+      else:
+        cell = SearchCell(C_prev, C_curr, 1, max_nodes, search_space)
+        if num_edge is None: num_edge, edge2index = cell.num_edges, cell.edge2index
+        else: assert num_edge == cell.num_edges and edge2index == cell.edge2index, 'invalid {:} vs. {:}.'.format(num_edge, cell.num_edges)
+      self.cells.append( cell )
+      C_prev = cell.out_dim
+    self.op_names   = deepcopy( search_space )
+    self._Layer     = len(self.cells)
+    self.edge2index = edge2index
+    self.lastact    = nn.Sequential(nn.BatchNorm2d(C_prev), nn.ReLU(inplace=True))
+    self.global_pooling = nn.AdaptiveAvgPool2d(1)
+    self.classifier = nn.Linear(C_prev, num_classes)
+    self.arch_parameters = nn.Parameter( 1e-3*torch.randn(num_edge, len(search_space)) )
+    self.mode       = 'urs'
+    self.dynamic_cell = None
+    
+  def set_cal_mode(self, mode, dynamic_cell=None):
+    assert mode in ['urs', 'joint', 'select', 'dynamic']
+    self.mode = mode
+    if mode == 'dynamic': self.dynamic_cell = deepcopy( dynamic_cell )
+    else                : self.dynamic_cell = None
+
+  def get_cal_mode(self):
+    return self.mode
+
+  def get_weights(self):
+    xlist = list( self.stem.parameters() ) + list( self.cells.parameters() )
+    xlist+= list( self.lastact.parameters() ) + list( self.global_pooling.parameters() )
+    xlist+= list( self.classifier.parameters() )
+    return xlist
+
+  def get_alphas(self):
+    return [self.arch_parameters]
+
+  def get_message(self):
+    string = self.extra_repr()
+    for i, cell in enumerate(self.cells):
+      string += '\n {:02d}/{:02d} :: {:}'.format(i, len(self.cells), cell.extra_repr())
+    return string
+
+  def extra_repr(self):
+    return ('{name}(C={_C}, Max-Nodes={max_nodes}, N={_layerN}, L={_Layer})'.format(name=self.__class__.__name__, **self.__dict__))
+
+  def genotype(self):
+    genotypes = []
+    for i in range(1, self.max_nodes):
+      xlist = []
+      for j in range(i):
+        node_str = '{:}<-{:}'.format(i, j)
+        with torch.no_grad():
+          weights = self.arch_parameters[ self.edge2index[node_str] ]
+          op_name = self.op_names[ weights.argmax().item() ]
+        xlist.append((op_name, j))
+      genotypes.append( tuple(xlist) )
+    return Structure( genotypes )
+
+
+  def dync_genotype(self):
+    genotypes = []
+    with torch.no_grad():
+      alphas_cpu = nn.functional.softmax(self.arch_parameters, dim=-1)
+    for i in range(1, self.max_nodes):
+      xlist = []
+      for j in range(i):
+        node_str = '{:}<-{:}'.format(i, j)
+        weights  = alphas_cpu[ self.edge2index[node_str] ]
+        op_index = torch.multinomial(weights, 1).item()
+        op_name  = self.op_names[ op_index ]
+        xlist.append((op_name, j))
+      genotypes.append( tuple(xlist) )
+    return Structure( genotypes )
+
+
+  def forward(self, inputs):
+    alphas  = nn.functional.softmax(self.arch_parameters, dim=-1)
+    with torch.no_grad():
+      alphas_cpu = alphas.detach().cpu()
+
+    feature = self.stem(inputs)
+    for i, cell in enumerate(self.cells):
+      if isinstance(cell, SearchCell):
+        if self.mode == 'urs':
+          feature = cell.forward_urs(feature)
+        elif self.mode == 'select':
+          feature = cell.forward_select(feature, alphas_cpu)
+        elif self.mode == 'joint':
+          feature = cell.forward_joint(feature, alphas)
+        elif self.mode == 'dynamic':
+          feature = cell.forward_dynamic(feature, self.dynamic_cell)
+        else: raise ValueError('invalid mode={:}'.format(self.mode))
+      else: feature = cell(feature)
+
+    out = self.lastact(feature)
+    out = self.global_pooling( out )
+    out = out.view(out.size(0), -1)
+    logits = self.classifier(out)
+
+    return out, logits