Update TAS abd FBV2 for NAS-Bench

lib/models/shape_searchs/__init__.py

@@ -6,3 +6,4 @@ from .SearchCifarResNet_depth import SearchDepthCifarResNet
 from .SearchCifarResNet       import SearchShapeCifarResNet
 from .SearchSimResNet_width   import SearchWidthSimResNet
 from .SearchImagenetResNet    import SearchShapeImagenetResNet
+from .generic_size_tiny_cell_model import GenericNAS301Model

lib/models/shape_searchs/generic_size_tiny_cell_model.py

@@ -0,0 +1,139 @@
+#####################################################
+# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2019.01 #
+#####################################################
+from typing import List, Text, Any
+import random, torch
+import torch.nn as nn
+
+from models.cell_operations import ResNetBasicblock
+from models.cell_infers.cells import InferCell
+from models.shape_searchs.SoftSelect import select2withP, ChannelWiseInter
+
+
+class GenericNAS301Model(nn.Module):
+
+  def __init__(self, candidate_Cs: List[int], max_num_Cs: int, genotype: Any, num_classes: int, affine: bool, track_running_stats: bool):
+    super(GenericNAS301Model, self).__init__()
+    self._max_num_Cs = max_num_Cs
+    self._candidate_Cs = candidate_Cs
+    if max_num_Cs % 3 != 2:
+      raise ValueError('invalid number of layers : {:}'.format(max_num_Cs))
+    self._num_stage = N = max_num_Cs // 3
+    self._max_C = max(candidate_Cs)
+
+    stem = nn.Sequential(
+                    nn.Conv2d(3, self._max_C, kernel_size=3, padding=1, bias=not affine),
+                    nn.BatchNorm2d(self._max_C, affine=affine, track_running_stats=track_running_stats))
+
+    layer_reductions = [False] * N + [True] + [False] * N + [True] + [False] * N
+
+    c_prev = self._max_C
+    self._cells = nn.ModuleList()
+    self._cells.append(stem)
+    for index, reduction in enumerate(layer_reductions):
+      if reduction : cell = ResNetBasicblock(c_prev, self._max_C, 2, True)
+      else         : cell = InferCell(genotype, c_prev, self._max_C, 1, affine, track_running_stats)
+      self._cells.append(cell)
+      c_prev = cell.out_dim
+    self._num_layer = len(self._cells)
+
+    self.lastact = nn.Sequential(nn.BatchNorm2d(c_prev, affine=affine, track_running_stats=track_running_stats), nn.ReLU(inplace=True))
+    self.global_pooling = nn.AdaptiveAvgPool2d(1)
+    self.classifier = nn.Linear(c_prev, num_classes)
+    # algorithm related
+    self.register_buffer('_tau', torch.zeros(1))
+    self._algo        = None
+
+  def set_algo(self, algo: Text):
+    # used for searching
+    assert self._algo is None, 'This functioin can only be called once.'
+    assert algo in ['fbv2', 'enas', 'tas'], 'invalid algo : {:}'.format(algo)
+    self._algo = algo
+    self._arch_parameters = nn.Parameter(1e-3*torch.randn(self._max_num_Cs, len(self._candidate_Cs)))
+    if algo == 'fbv2' or algo == 'enas':
+      self.register_buffer('_masks', torch.zeros(len(self._candidate_Cs), max(self._candidate_Cs)))
+      for i in range(len(self._candidate_Cs)):
+        self._masks.data[i, :self._candidate_Cs[i]] = 1
+  
+  @property
+  def tau(self):
+    return self._tau
+
+  def set_tau(self, tau):
+    self._tau.data[:] = tau
+
+  @property
+  def weights(self):
+    xlist = list(self._cells.parameters())
+    xlist+= list(self.lastact.parameters())
+    xlist+= list(self.global_pooling.parameters())
+    xlist+= list(self.classifier.parameters())
+    return xlist
+
+  @property
+  def alphas(self):
+    return [self._arch_parameters]
+
+  def show_alphas(self):
+    with torch.no_grad():
+      return 'arch-parameters :\n{:}'.format(nn.functional.softmax(self._arch_parameters, dim=-1).cpu())
+
+  @property
+  def random(self):
+    cs = []
+    for i in range(self._max_num_Cs):
+      index = random.randint(0, len(self._candidate_Cs)-1)
+      cs.append(str(self._candidate_Cs[index]))
+    return ':'.join(cs)
+  
+  @property
+  def genotype(self):
+    cs = []
+    for i in range(self._max_num_Cs):
+      with torch.no_grad():
+        index = self._arch_parameters[i].argmax().item()
+        cs.append(str(self._candidate_Cs[index]))
+    return ':'.join(cs)
+
+  def get_message(self) -> Text:
+    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}(candidates={_candidate_Cs}, num={_max_num_Cs}, N={_num_stage}, L={_num_layer})'.format(name=self.__class__.__name__, **self.__dict__))
+
+  def forward(self, inputs):
+    feature = inputs
+    for i, cell in enumerate(self._cells):
+      feature = cell(feature)
+      if self._algo == 'fbv2':
+        idx = max(0, i-1)
+        weights = nn.functional.gumbel_softmax(self._arch_parameters[idx:idx+1], tau=self.tau, dim=-1)
+        mask = torch.matmul(weights, self._masks).view(1, -1, 1, 1)
+        feature = feature * mask
+      elif self._algo == 'tas':
+        idx = max(0, i-1)
+        selected_cs, selected_probs = select2withP(self._arch_parameters[idx:idx+1], self.tau, num=2)
+        with torch.no_grad():
+          i1, i2 = selected_cs.cpu().view(-1).tolist()
+        c1, c2 = self._candidate_Cs[i1], self._candidate_Cs[i2]
+        out_channel = max(c1, c2)
+        out1 = ChannelWiseInter(feature[:, :c1], out_channel)
+        out2 = ChannelWiseInter(feature[:, :c2], out_channel)
+        out  = out1 * selected_probs[0, 0] + out2 * selected_probs[0, 1]
+        if feature.shape[1] == out.shape[1]:
+          feature = out
+        else:
+          miss = torch.zeros(feature.shape[0], feature.shape[1]-out.shape[1], feature.shape[2], feature.shape[3], device=feature.device)
+          feature = torch.cat((out, miss), dim=1)
+      else:
+        raise ValueError('invalid algorithm : {:}'.format(self._algo))
+
+    out = self.lastact(feature)
+    out = self.global_pooling(out)
+    out = out.view(out.size(0), -1)
+    logits = self.classifier(out)
+
+    return out, logits