Add more algorithms

lib/models/CifarDenseNet.py

@@ -0,0 +1,105 @@
+##################################################
+# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2019 #
+##################################################
+import math, torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .initialization import initialize_resnet
+
+
+class Bottleneck(nn.Module):
+  def __init__(self, nChannels, growthRate):
+    super(Bottleneck, self).__init__()
+    interChannels = 4*growthRate
+    self.bn1 = nn.BatchNorm2d(nChannels)
+    self.conv1 = nn.Conv2d(nChannels, interChannels, kernel_size=1, bias=False)
+    self.bn2 = nn.BatchNorm2d(interChannels)
+    self.conv2 = nn.Conv2d(interChannels, growthRate, kernel_size=3, padding=1, bias=False)
+
+  def forward(self, x):
+    out = self.conv1(F.relu(self.bn1(x)))
+    out = self.conv2(F.relu(self.bn2(out)))
+    out = torch.cat((x, out), 1)
+    return out
+
+
+class SingleLayer(nn.Module):
+  def __init__(self, nChannels, growthRate):
+    super(SingleLayer, self).__init__()
+    self.bn1 = nn.BatchNorm2d(nChannels)
+    self.conv1 = nn.Conv2d(nChannels, growthRate, kernel_size=3, padding=1, bias=False)
+
+  def forward(self, x):
+    out = self.conv1(F.relu(self.bn1(x)))
+    out = torch.cat((x, out), 1)
+    return out
+
+
+class Transition(nn.Module):
+  def __init__(self, nChannels, nOutChannels):
+    super(Transition, self).__init__()
+    self.bn1 = nn.BatchNorm2d(nChannels)
+    self.conv1 = nn.Conv2d(nChannels, nOutChannels, kernel_size=1, bias=False)
+
+  def forward(self, x):
+    out = self.conv1(F.relu(self.bn1(x)))
+    out = F.avg_pool2d(out, 2)
+    return out
+
+
+class DenseNet(nn.Module):
+  def __init__(self, growthRate, depth, reduction, nClasses, bottleneck):
+    super(DenseNet, self).__init__()
+
+    if bottleneck:  nDenseBlocks = int( (depth-4) / 6 )
+    else         :  nDenseBlocks = int( (depth-4) / 3 )
+
+    self.message = 'CifarDenseNet : block : {:}, depth : {:}, reduction : {:}, growth-rate = {:}, class = {:}'.format('bottleneck' if bottleneck else 'basic', depth, reduction, growthRate, nClasses)
+
+    nChannels = 2*growthRate
+    self.conv1 = nn.Conv2d(3, nChannels, kernel_size=3, padding=1, bias=False)
+
+    self.dense1 = self._make_dense(nChannels, growthRate, nDenseBlocks, bottleneck)
+    nChannels += nDenseBlocks*growthRate
+    nOutChannels = int(math.floor(nChannels*reduction))
+    self.trans1 = Transition(nChannels, nOutChannels)
+
+    nChannels = nOutChannels
+    self.dense2 = self._make_dense(nChannels, growthRate, nDenseBlocks, bottleneck)
+    nChannels += nDenseBlocks*growthRate
+    nOutChannels = int(math.floor(nChannels*reduction))
+    self.trans2 = Transition(nChannels, nOutChannels)
+
+    nChannels = nOutChannels
+    self.dense3 = self._make_dense(nChannels, growthRate, nDenseBlocks, bottleneck)
+    nChannels += nDenseBlocks*growthRate
+
+    self.act = nn.Sequential(
+                  nn.BatchNorm2d(nChannels), nn.ReLU(inplace=True),
+                  nn.AvgPool2d(8))
+    self.fc  = nn.Linear(nChannels, nClasses)
+
+    self.apply(initialize_resnet)
+
+  def get_message(self):
+    return self.message
+
+  def _make_dense(self, nChannels, growthRate, nDenseBlocks, bottleneck):
+    layers = []
+    for i in range(int(nDenseBlocks)):
+      if bottleneck:
+        layers.append(Bottleneck(nChannels, growthRate))
+      else:
+        layers.append(SingleLayer(nChannels, growthRate))
+      nChannels += growthRate
+    return nn.Sequential(*layers)
+
+  def forward(self, inputs):
+    out = self.conv1( inputs )
+    out = self.trans1(self.dense1(out))
+    out = self.trans2(self.dense2(out))
+    out = self.dense3(out)
+    features = self.act(out)
+    features = features.view(features.size(0), -1)
+    out = self.fc(features)
+    return features, out

lib/models/CifarResNet.py

@@ -0,0 +1,160 @@
+##################################################
+# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2019 #
+##################################################
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .initialization import initialize_resnet
+from .SharedUtils    import additive_func
+
+
+class Downsample(nn.Module):  
+
+  def __init__(self, nIn, nOut, stride):
+    super(Downsample, self).__init__() 
+    assert stride == 2 and nOut == 2*nIn, 'stride:{} IO:{},{}'.format(stride, nIn, nOut)
+    self.in_dim  = nIn
+    self.out_dim = nOut
+    self.avg  = nn.AvgPool2d(kernel_size=2, stride=2, padding=0)   
+    self.conv = nn.Conv2d(nIn, nOut, kernel_size=1, stride=1, padding=0, bias=False)
+
+  def forward(self, x):
+    x   = self.avg(x)
+    out = self.conv(x)
+    return out
+
+
+class ConvBNReLU(nn.Module):
+  
+  def __init__(self, nIn, nOut, kernel, stride, padding, bias, relu):
+    super(ConvBNReLU, self).__init__()
+    self.conv = nn.Conv2d(nIn, nOut, kernel_size=kernel, stride=stride, padding=padding, bias=bias)
+    self.bn   = nn.BatchNorm2d(nOut)
+    if relu: self.relu = nn.ReLU(inplace=True)
+    else   : self.relu = None
+    self.out_dim = nOut
+    self.num_conv = 1
+
+  def forward(self, x):
+    conv = self.conv( x )
+    bn   = self.bn( conv )
+    if self.relu: return self.relu( bn )
+    else        : return bn
+
+
+class ResNetBasicblock(nn.Module):
+  expansion = 1
+  def __init__(self, inplanes, planes, stride):
+    super(ResNetBasicblock, self).__init__()
+    assert stride == 1 or stride == 2, 'invalid stride {:}'.format(stride)
+    self.conv_a = ConvBNReLU(inplanes, planes, 3, stride, 1, False, True)
+    self.conv_b = ConvBNReLU(  planes, planes, 3,      1, 1, False, False)
+    if stride == 2:
+      self.downsample = Downsample(inplanes, planes, stride)
+    elif inplanes != planes:
+      self.downsample = ConvBNReLU(inplanes, planes, 1, 1, 0, False, False)
+    else:
+      self.downsample = None
+    self.out_dim = planes
+    self.num_conv = 2
+
+  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
+    out = additive_func(residual, basicblock)
+    return F.relu(out, inplace=True)
+
+
+
+class ResNetBottleneck(nn.Module):
+  expansion = 4
+  def __init__(self, inplanes, planes, stride):
+    super(ResNetBottleneck, self).__init__()
+    assert stride == 1 or stride == 2, 'invalid stride {:}'.format(stride)
+    self.conv_1x1 = ConvBNReLU(inplanes, planes, 1,      1, 0, False, True)
+    self.conv_3x3 = ConvBNReLU(  planes, planes, 3, stride, 1, False, True)
+    self.conv_1x4 = ConvBNReLU(planes, planes*self.expansion, 1, 1, 0, False, False)
+    if stride == 2:
+      self.downsample = Downsample(inplanes, planes*self.expansion, stride)
+    elif inplanes != planes*self.expansion:
+      self.downsample = ConvBNReLU(inplanes, planes*self.expansion, 1, 1, 0, False, False)
+    else:
+      self.downsample = None
+    self.out_dim = planes * self.expansion
+    self.num_conv = 3
+
+  def forward(self, inputs):
+
+    bottleneck = self.conv_1x1(inputs)
+    bottleneck = self.conv_3x3(bottleneck)
+    bottleneck = self.conv_1x4(bottleneck)
+
+    if self.downsample is not None:
+      residual = self.downsample(inputs)
+    else:
+      residual = inputs
+    out = additive_func(residual, bottleneck)
+    return F.relu(out, inplace=True)
+
+
+
+class CifarResNet(nn.Module):
+
+  def __init__(self, block_name, depth, num_classes, zero_init_residual):
+    super(CifarResNet, self).__init__()
+
+    #Model type specifies number of layers for CIFAR-10 and CIFAR-100 model
+    if block_name == 'ResNetBasicblock':
+      block = ResNetBasicblock
+      assert (depth - 2) % 6 == 0, 'depth should be one of 20, 32, 44, 56, 110'
+      layer_blocks = (depth - 2) // 6
+    elif block_name == 'ResNetBottleneck':
+      block = ResNetBottleneck
+      assert (depth - 2) % 9 == 0, 'depth should be one of 164'
+      layer_blocks = (depth - 2) // 9
+    else:
+      raise ValueError('invalid block : {:}'.format(block_name))
+
+    self.message     = 'CifarResNet : Block : {:}, Depth : {:}, Layers for each block : {:}'.format(block_name, depth, layer_blocks)
+    self.num_classes = num_classes
+    self.channels    = [16]
+    self.layers      = nn.ModuleList( [ ConvBNReLU(3, 16, 3, 1, 1, False, True) ] )
+    for stage in range(3):
+      for iL in range(layer_blocks):
+        iC     = self.channels[-1]
+        planes = 16 * (2**stage)
+        stride = 2 if stage > 0 and iL == 0 else 1
+        module = block(iC, planes, stride)
+        self.channels.append( module.out_dim )
+        self.layers.append  ( module )
+        self.message += "\nstage={:}, ilayer={:02d}/{:02d}, block={:03d}, iC={:3d}, oC={:3d}, stride={:}".format(stage, iL, layer_blocks, len(self.layers)-1, iC, module.out_dim, stride)
+
+    self.avgpool = nn.AvgPool2d(8)
+    self.classifier = nn.Linear(module.out_dim, num_classes)
+    assert sum(x.num_conv for x in self.layers) + 1 == depth, 'invalid depth check {:} vs {:}'.format(sum(x.num_conv for x in self.layers)+1, depth)
+
+    self.apply(initialize_resnet)
+    if zero_init_residual:
+      for m in self.modules():
+        if isinstance(m, ResNetBasicblock):
+          nn.init.constant_(m.conv_b.bn.weight, 0)
+        elif isinstance(m, ResNetBottleneck):
+          nn.init.constant_(m.conv_1x4.bn.weight, 0)
+
+  def get_message(self):
+    return self.message
+
+  def forward(self, inputs):
+    x = inputs
+    for i, layer in enumerate(self.layers):
+      x = layer( x )
+    features = self.avgpool(x)
+    features = features.view(features.size(0), -1)
+    logits   = self.classifier(features)
+    return features, logits

lib/models/CifarWideResNet.py

@@ -0,0 +1,94 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .initialization import initialize_resnet
+
+
+class WideBasicblock(nn.Module):
+  def __init__(self, inplanes, planes, stride, dropout=False):
+    super(WideBasicblock, self).__init__()
+
+    self.bn_a = nn.BatchNorm2d(inplanes)
+    self.conv_a = nn.Conv2d(inplanes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
+
+    self.bn_b = nn.BatchNorm2d(planes)
+    if dropout:
+      self.dropout = nn.Dropout2d(p=0.5, inplace=True)
+    else:
+      self.dropout = None
+    self.conv_b = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
+
+    if inplanes != planes:
+      self.downsample = nn.Conv2d(inplanes, planes, kernel_size=1, stride=stride, padding=0, bias=False)
+    else:
+      self.downsample = None
+
+  def forward(self, x):
+
+    basicblock = self.bn_a(x)
+    basicblock = F.relu(basicblock)
+    basicblock = self.conv_a(basicblock)
+
+    basicblock = self.bn_b(basicblock)
+    basicblock = F.relu(basicblock)
+    if self.dropout is not None:
+      basicblock = self.dropout(basicblock)
+    basicblock = self.conv_b(basicblock)
+
+    if self.downsample is not None:
+      x = self.downsample(x)
+    
+    return x + basicblock
+
+
+class CifarWideResNet(nn.Module):
+  """
+  ResNet optimized for the Cifar dataset, as specified in
+  https://arxiv.org/abs/1512.03385.pdf
+  """
+  def __init__(self, depth, widen_factor, num_classes, dropout):
+    super(CifarWideResNet, self).__init__()
+
+    #Model type specifies number of layers for CIFAR-10 and CIFAR-100 model
+    assert (depth - 4) % 6 == 0, 'depth should be one of 20, 32, 44, 56, 110'
+    layer_blocks = (depth - 4) // 6
+    print ('CifarPreResNet : Depth : {} , Layers for each block : {}'.format(depth, layer_blocks))
+
+    self.num_classes = num_classes
+    self.dropout = dropout
+    self.conv_3x3 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1, bias=False)
+
+    self.message  = 'Wide ResNet : depth={:}, widen_factor={:}, class={:}'.format(depth, widen_factor, num_classes)
+    self.inplanes = 16
+    self.stage_1 = self._make_layer(WideBasicblock, 16*widen_factor, layer_blocks, 1)
+    self.stage_2 = self._make_layer(WideBasicblock, 32*widen_factor, layer_blocks, 2)
+    self.stage_3 = self._make_layer(WideBasicblock, 64*widen_factor, layer_blocks, 2)
+    self.lastact = nn.Sequential(nn.BatchNorm2d(64*widen_factor), nn.ReLU(inplace=True))
+    self.avgpool = nn.AvgPool2d(8)
+    self.classifier = nn.Linear(64*widen_factor, num_classes)
+
+    self.apply(initialize_resnet)
+
+  def get_message(self):
+    return self.message
+
+  def _make_layer(self, block, planes, blocks, stride):
+
+    layers = []
+    layers.append(block(self.inplanes, planes, stride, self.dropout))
+    self.inplanes = planes
+    for i in range(1, blocks):
+      layers.append(block(self.inplanes, planes, 1, self.dropout))
+
+    return nn.Sequential(*layers)
+
+  def forward(self, x):
+    x = self.conv_3x3(x)
+    x = self.stage_1(x)
+    x = self.stage_2(x)
+    x = self.stage_3(x)
+    x = self.lastact(x)
+    x = self.avgpool(x)
+    features = x.view(x.size(0), -1)
+    outs     = self.classifier(features)
+    return features, outs

lib/models/ImagenetResNet.py

@@ -0,0 +1,172 @@
+# Deep Residual Learning for Image Recognition, CVPR 2016
+import torch.nn as nn
+from .initialization import initialize_resnet
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1):
+  return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, groups=groups, bias=False)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+  return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class BasicBlock(nn.Module):
+  expansion = 1
+
+  def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, base_width=64):
+    super(BasicBlock, self).__init__()
+    if groups != 1 or base_width != 64:
+      raise ValueError('BasicBlock only supports groups=1 and base_width=64')
+    # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+    self.conv1 = conv3x3(inplanes, planes, stride)
+    self.bn1   = nn.BatchNorm2d(planes)
+    self.relu  = nn.ReLU(inplace=True)
+    self.conv2 = conv3x3(planes, planes)
+    self.bn2   = nn.BatchNorm2d(planes)
+    self.downsample = downsample
+    self.stride = stride
+
+  def forward(self, x):
+    identity = x
+
+    out = self.conv1(x)
+    out = self.bn1(out)
+    out = self.relu(out)
+
+    out = self.conv2(out)
+    out = self.bn2(out)
+
+    if self.downsample is not None:
+      identity = self.downsample(x)
+
+    out += identity
+    out = self.relu(out)
+
+    return out
+
+
+class Bottleneck(nn.Module):
+  expansion = 4
+
+  def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, base_width=64):
+    super(Bottleneck, self).__init__()
+    width = int(planes * (base_width / 64.)) * groups
+    # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+    self.conv1 = conv1x1(inplanes, width)
+    self.bn1   = nn.BatchNorm2d(width)
+    self.conv2 = conv3x3(width, width, stride, groups)
+    self.bn2   = nn.BatchNorm2d(width)
+    self.conv3 = conv1x1(width, planes * self.expansion)
+    self.bn3   = nn.BatchNorm2d(planes * self.expansion)
+    self.relu  = nn.ReLU(inplace=True)
+    self.downsample = downsample
+    self.stride = stride
+
+  def forward(self, x):
+    identity = x
+
+    out = self.conv1(x)
+    out = self.bn1(out)
+    out = self.relu(out)
+
+    out = self.conv2(out)
+    out = self.bn2(out)
+    out = self.relu(out)
+
+    out = self.conv3(out)
+    out = self.bn3(out)
+
+    if self.downsample is not None:
+      identity = self.downsample(x)
+
+    out += identity
+    out = self.relu(out)
+
+    return out
+
+
+class ResNet(nn.Module):
+
+  def __init__(self, block_name, layers, deep_stem, num_classes, zero_init_residual, groups, width_per_group):
+    super(ResNet, self).__init__()
+
+    #planes = [int(width_per_group * groups * 2 ** i) for i in range(4)]
+    if block_name == 'BasicBlock'  : block= BasicBlock
+    elif block_name == 'Bottleneck': block= Bottleneck
+    else                           : raise ValueError('invalid block-name : {:}'.format(block_name))
+
+    if not deep_stem:
+      self.conv = nn.Sequential(
+                   nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False),
+                   nn.BatchNorm2d(64), nn.ReLU(inplace=True))
+    else:
+      self.conv = nn.Sequential(
+                   nn.Conv2d(           3, 32, kernel_size=3, stride=2, padding=1, bias=False),
+                   nn.BatchNorm2d(32), nn.ReLU(inplace=True),
+                   nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1, bias=False),
+                   nn.BatchNorm2d(32), nn.ReLU(inplace=True),
+                   nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1, bias=False),
+                   nn.BatchNorm2d(64), nn.ReLU(inplace=True))
+    self.inplanes = 64
+    self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+    self.layer1 = self._make_layer(block, 64 , layers[0], stride=1, groups=groups, base_width=width_per_group)
+    self.layer2 = self._make_layer(block, 128, layers[1], stride=2, groups=groups, base_width=width_per_group)
+    self.layer3 = self._make_layer(block, 256, layers[2], stride=2, groups=groups, base_width=width_per_group)
+    self.layer4 = self._make_layer(block, 512, layers[3], stride=2, groups=groups, base_width=width_per_group)
+    self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+    self.fc      = nn.Linear(512 * block.expansion, num_classes)
+    self.message = 'block = {:}, layers = {:}, deep_stem = {:}, num_classes = {:}'.format(block, layers, deep_stem, num_classes)
+
+    self.apply( initialize_resnet )
+
+    # Zero-initialize the last BN in each residual branch,
+    # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+    # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+    if zero_init_residual:
+      for m in self.modules():
+        if isinstance(m, Bottleneck):
+          nn.init.constant_(m.bn3.weight, 0)
+        elif isinstance(m, BasicBlock):
+          nn.init.constant_(m.bn2.weight, 0)
+
+  def _make_layer(self, block, planes, blocks, stride, groups, base_width):
+    downsample = None
+    if stride != 1 or self.inplanes != planes * block.expansion:
+      if stride == 2:
+        downsample = nn.Sequential(
+          nn.AvgPool2d(kernel_size=2, stride=2, padding=0),
+          conv1x1(self.inplanes, planes * block.expansion, 1),
+          nn.BatchNorm2d(planes * block.expansion),
+        )
+      elif stride == 1:
+        downsample = nn.Sequential(
+          conv1x1(self.inplanes, planes * block.expansion, stride),
+          nn.BatchNorm2d(planes * block.expansion),
+        )
+      else: raise ValueError('invalid stride [{:}] for downsample'.format(stride))
+
+    layers = []
+    layers.append(block(self.inplanes, planes, stride, downsample, groups, base_width))
+    self.inplanes = planes * block.expansion
+    for _ in range(1, blocks):
+      layers.append(block(self.inplanes, planes, 1, None, groups, base_width))
+
+    return nn.Sequential(*layers)
+
+  def get_message(self):
+    return self.message
+
+  def forward(self, x):
+    x = self.conv(x)
+    x = self.maxpool(x)
+
+    x = self.layer1(x)
+    x = self.layer2(x)
+    x = self.layer3(x)
+    x = self.layer4(x)
+
+    features = self.avgpool(x)
+    features = features.view(features.size(0), -1)
+    logits   = self.fc(features)
+
+    return features, logits

lib/models/MobileNet.py

@@ -0,0 +1,101 @@
+# MobileNetV2: Inverted Residuals and Linear Bottlenecks, CVPR 2018
+from torch import nn
+from .initialization import initialize_resnet
+
+
+class ConvBNReLU(nn.Module):
+  def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1):
+    super(ConvBNReLU, self).__init__()
+    padding = (kernel_size - 1) // 2
+    self.conv = nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False)
+    self.bn   = nn.BatchNorm2d(out_planes)
+    self.relu = nn.ReLU6(inplace=True)
+  
+  def forward(self, x):
+    out = self.conv( x )
+    out = self.bn  ( out )
+    out = self.relu( out )
+    return out
+
+
+class InvertedResidual(nn.Module):
+  def __init__(self, inp, oup, stride, expand_ratio):
+    super(InvertedResidual, self).__init__()
+    self.stride = stride
+    assert stride in [1, 2]
+
+    hidden_dim = int(round(inp * expand_ratio))
+    self.use_res_connect = self.stride == 1 and inp == oup
+
+    layers = []
+    if expand_ratio != 1:
+      # pw
+      layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1))
+    layers.extend([
+      # dw
+      ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim),
+      # pw-linear
+      nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+      nn.BatchNorm2d(oup),
+    ])
+    self.conv = nn.Sequential(*layers)
+
+  def forward(self, x):
+    if self.use_res_connect:
+      return x + self.conv(x)
+    else:
+      return self.conv(x)
+
+
+class MobileNetV2(nn.Module):
+  def __init__(self, num_classes, width_mult, input_channel, last_channel, block_name, dropout):
+    super(MobileNetV2, self).__init__()
+    if block_name == 'InvertedResidual':
+      block = InvertedResidual
+    else:
+      raise ValueError('invalid block name : {:}'.format(block_name))
+    inverted_residual_setting = [
+      # t, c,  n, s
+      [1, 16 , 1, 1],
+      [6, 24 , 2, 2],
+      [6, 32 , 3, 2],
+      [6, 64 , 4, 2],
+      [6, 96 , 3, 1],
+      [6, 160, 3, 2],
+      [6, 320, 1, 1],
+    ]
+
+    # building first layer
+    input_channel = int(input_channel * width_mult)
+    self.last_channel = int(last_channel * max(1.0, width_mult))
+    features = [ConvBNReLU(3, input_channel, stride=2)]
+    # building inverted residual blocks
+    for t, c, n, s in inverted_residual_setting:
+      output_channel = int(c * width_mult)
+      for i in range(n):
+        stride = s if i == 0 else 1
+        features.append(block(input_channel, output_channel, stride, expand_ratio=t))
+        input_channel = output_channel
+    # building last several layers
+    features.append(ConvBNReLU(input_channel, self.last_channel, kernel_size=1))
+    # make it nn.Sequential
+    self.features = nn.Sequential(*features)
+
+    # building classifier
+    self.classifier = nn.Sequential(
+      nn.Dropout(dropout),
+      nn.Linear(self.last_channel, num_classes),
+    )
+    self.message = 'MobileNetV2 : width_mult={:}, in-C={:}, last-C={:}, block={:}, dropout={:}'.format(width_mult, input_channel, last_channel, block_name, dropout)
+
+    # weight initialization
+    self.apply( initialize_resnet )
+
+  def get_message(self):
+    return self.message
+
+  def forward(self, inputs):
+    features = self.features(inputs)
+    vectors  = features.mean([2, 3])
+    predicts = self.classifier(vectors)
+    return features, predicts

lib/models/SharedUtils.py

@@ -0,0 +1,31 @@
+import torch
+import torch.nn as nn
+
+
+def additive_func(A, B):
+  assert A.dim() == B.dim() and A.size(0) == B.size(0), '{:} vs {:}'.format(A.size(), B.size())
+  C = min(A.size(1), B.size(1))
+  if A.size(1) == B.size(1):
+    return A + B
+  elif A.size(1) < B.size(1):
+    out = B.clone()
+    out[:,:C] += A
+    return out
+  else:
+    out = A.clone()
+    out[:,:C] += B
+    return out
+
+
+def change_key(key, value):
+  def func(m):
+    if hasattr(m, key):
+      setattr(m, key, value)
+  return func
+
+
+def parse_channel_info(xstring):
+  blocks = xstring.split(' ')
+  blocks = [x.split('-') for x in blocks]
+  blocks = [[int(_) for _ in x] for x in blocks]
+  return blocks

lib/models/ShuffleNetV2.py

@@ -0,0 +1,133 @@
+import functools
+
+import torch
+import torch.nn as nn
+
+__all__ = ['ShuffleNetV2']
+
+
+def channel_shuffle(x, groups):
+  batchsize, num_channels, height, width = x.data.size()
+  channels_per_group = num_channels // groups
+
+  # reshape
+  x = x.view(batchsize, groups, channels_per_group, height, width)
+
+  x = torch.transpose(x, 1, 2).contiguous()
+
+  # flatten
+  x = x.view(batchsize, -1, height, width)
+
+  return x
+
+
+class InvertedResidual(nn.Module):
+  def __init__(self, inp, oup, stride):
+    super(InvertedResidual, self).__init__()
+
+    if not (1 <= stride <= 3):
+      raise ValueError('illegal stride value')
+    self.stride = stride
+
+    branch_features = oup // 2
+    assert (self.stride != 1) or (inp == branch_features << 1)
+
+    pw_conv11 = functools.partial(nn.Conv2d, kernel_size=1, stride=1, padding=0, bias=False)
+    dw_conv33 = functools.partial(self.depthwise_conv, kernel_size=3, stride=self.stride, padding=1)
+
+    if self.stride > 1:
+      self.branch1 = nn.Sequential(
+        dw_conv33(inp, inp),
+        nn.BatchNorm2d(inp),
+        pw_conv11(inp, branch_features),
+        nn.BatchNorm2d(branch_features),
+        nn.ReLU(inplace=True),
+      )
+
+    self.branch2 = nn.Sequential(
+      pw_conv11(inp if (self.stride > 1) else branch_features, branch_features),
+      nn.BatchNorm2d(branch_features),
+      nn.ReLU(inplace=True),
+      dw_conv33(branch_features, branch_features),
+      nn.BatchNorm2d(branch_features),
+      pw_conv11(branch_features, branch_features),
+      nn.BatchNorm2d(branch_features),
+      nn.ReLU(inplace=True),
+    )
+
+  @staticmethod
+  def depthwise_conv(i, o, kernel_size, stride=1, padding=0, bias=False):
+    return nn.Conv2d(i, o, kernel_size, stride, padding, bias=bias, groups=i)
+
+  def forward(self, x):
+    if self.stride == 1:
+      x1, x2 = x.chunk(2, dim=1)
+      out = torch.cat((x1, self.branch2(x2)), dim=1)
+    else:
+      out = torch.cat((self.branch1(x), self.branch2(x)), dim=1)
+
+    out = channel_shuffle(out, 2)
+    return out
+
+
+class ShuffleNetV2(nn.Module):
+  def __init__(self, num_classes, stages):
+    super(ShuffleNetV2, self).__init__()
+
+    self.stage_out_channels = stages
+    assert len(stages) == 5, 'invalid stages : {:}'.format(stages)
+    self.message = 'stages: ' + ' '.join([str(x) for x in stages])
+
+    input_channels = 3
+    output_channels = self.stage_out_channels[0]
+    self.conv1 = nn.Sequential(
+      nn.Conv2d(input_channels, output_channels, 3, 2, 1, bias=False),
+      nn.BatchNorm2d(output_channels),
+      nn.ReLU(inplace=True),
+    )
+    input_channels = output_channels
+
+    self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+    stage_names = ['stage{:}'.format(i) for i in [2, 3, 4]]
+    stage_repeats = [4, 8, 4]
+    for name, repeats, output_channels in zip(
+        stage_names, stage_repeats, self.stage_out_channels[1:]):
+      seq = [InvertedResidual(input_channels, output_channels, 2)]
+      for i in range(repeats - 1):
+        seq.append(InvertedResidual(output_channels, output_channels, 1))
+      setattr(self, name, nn.Sequential(*seq))
+      input_channels = output_channels
+
+    output_channels = self.stage_out_channels[-1]
+    self.conv5 = nn.Sequential(
+      nn.Conv2d(input_channels, output_channels, 1, 1, 0, bias=False),
+      nn.BatchNorm2d(output_channels),
+      nn.ReLU(inplace=True),
+    )
+
+    self.fc = nn.Linear(output_channels, num_classes)
+
+  def get_message(self):
+    return self.message
+
+  def forward(self, inputs):
+    x = self.conv1( inputs )
+    x = self.maxpool(x)
+    x = self.stage2(x)
+    x = self.stage3(x)
+    x = self.stage4(x)
+    x = self.conv5(x)
+    features = x.mean([2, 3])  # globalpool
+    predicts = self.fc(features)
+    return features, predicts
+
+  #@staticmethod
+  #def _getStages(mult):
+  #  stages = {
+  #    '0.5': [24, 48,  96 , 192, 1024],
+  #    '1.0': [24, 116, 232, 464, 1024],
+  #    '1.5': [24, 176, 352, 704, 1024],
+  #    '2.0': [24, 244, 488, 976, 2048],
+  #  }
+  #  return stages[str(mult)]

lib/models/__init__.py

@@ -0,0 +1,123 @@
+##################################################
+# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2019 #
+##################################################
+import torch
+from os import path as osp
+# our modules
+from config_utils import dict2config
+from .SharedUtils import change_key
+from .clone_weights import init_from_model
+
+
+def get_cifar_models(config):
+  from .CifarResNet      import CifarResNet
+  from .CifarDenseNet    import DenseNet
+  from .CifarWideResNet  import CifarWideResNet
+  
+  super_type = getattr(config, 'super_type', 'basic')
+  if super_type == 'basic':
+    if config.arch == 'resnet':
+      return CifarResNet(config.module, config.depth, config.class_num, config.zero_init_residual)
+    elif config.arch == 'densenet':
+      return DenseNet(config.growthRate, config.depth, config.reduction, config.class_num, config.bottleneck)
+    elif config.arch == 'wideresnet':
+      return CifarWideResNet(config.depth, config.wide_factor, config.class_num, config.dropout)
+    else:
+      raise ValueError('invalid module type : {:}'.format(config.arch))
+  elif super_type.startswith('infer'):
+    from .infers import InferWidthCifarResNet
+    from .infers import InferDepthCifarResNet
+    from .infers import InferCifarResNet
+    assert len(super_type.split('-')) == 2, 'invalid super_type : {:}'.format(super_type)
+    infer_mode = super_type.split('-')[1]
+    if infer_mode == 'width':
+      return InferWidthCifarResNet(config.module, config.depth, config.xchannels, config.class_num, config.zero_init_residual)
+    elif infer_mode == 'depth':
+      return InferDepthCifarResNet(config.module, config.depth, config.xblocks, config.class_num, config.zero_init_residual)
+    elif infer_mode == 'shape':
+      return InferCifarResNet(config.module, config.depth, config.xblocks, config.xchannels, config.class_num, config.zero_init_residual)
+    else:
+      raise ValueError('invalid infer-mode : {:}'.format(infer_mode))
+  else:
+    raise ValueError('invalid super-type : {:}'.format(super_type))
+
+
+def get_imagenet_models(config):
+  super_type = getattr(config, 'super_type', 'basic')
+  if super_type == 'basic':
+    return get_imagenet_models_basic(config)
+  # NAS searched architecture
+  elif super_type.startswith('infer'):
+    assert len(super_type.split('-')) == 2, 'invalid super_type : {:}'.format(super_type)
+    infer_mode = super_type.split('-')[1]
+    if infer_mode == 'shape':
+      from .infers import InferImagenetResNet
+      from .infers import InferMobileNetV2
+      if config.arch == 'resnet':
+        return InferImagenetResNet(config.block_name, config.layers, config.xblocks, config.xchannels, config.deep_stem, config.class_num, config.zero_init_residual)
+      elif config.arch == "MobileNetV2":
+        return InferMobileNetV2(config.class_num, config.xchannels, config.xblocks, config.dropout)
+      else:
+        raise ValueError('invalid arch-mode : {:}'.format(config.arch))
+    else:
+      raise ValueError('invalid infer-mode : {:}'.format(infer_mode))
+  else:
+    raise ValueError('invalid super-type : {:}'.format(super_type))
+
+
+def get_imagenet_models_basic(config):
+  from .ImagenetResNet import ResNet
+  from .MobileNet      import MobileNetV2
+  from .ShuffleNetV2   import ShuffleNetV2
+  if config.arch == 'resnet':
+    return ResNet(config.block_name, config.layers, config.deep_stem, config.class_num, config.zero_init_residual, config.groups, config.width_per_group)
+  elif config.arch == 'MobileNetV2':
+    return MobileNetV2(config.class_num, config.width_mult, config.input_channel, config.last_channel, config.block_name, config.dropout)
+  elif config.arch == 'ShuffleNetV2':
+    return ShuffleNetV2(config.class_num, config.stages)
+  else:
+    raise ValueError('invalid arch : {:}'.format( config.arch ))
+    
+
+def obtain_model(config):
+  if config.dataset == 'cifar':
+    return get_cifar_models(config)
+  elif config.dataset == 'imagenet':
+    return get_imagenet_models(config)
+  else:
+    raise ValueError('invalid dataset in the model config : {:}'.format(config))
+
+
+def obtain_search_model(config):
+  if config.dataset == 'cifar':
+    if config.arch == 'resnet':
+      from .searchs import SearchWidthCifarResNet
+      from .searchs import SearchDepthCifarResNet
+      from .searchs import SearchShapeCifarResNet
+      if config.search_mode == 'width':
+        return SearchWidthCifarResNet(config.module, config.depth, config.class_num)
+      elif config.search_mode == 'depth':
+        return SearchDepthCifarResNet(config.module, config.depth, config.class_num)
+      elif config.search_mode == 'shape':
+        return SearchShapeCifarResNet(config.module, config.depth, config.class_num)
+      else: raise ValueError('invalid search mode : {:}'.format(config.search_mode))
+    else:
+      raise ValueError('invalid arch : {:} for dataset [{:}]'.format(config.arch, config.dataset))
+  elif config.dataset == 'imagenet':
+    from .searchs import SearchShapeImagenetResNet
+    assert config.search_mode == 'shape', 'invalid search-mode : {:}'.format( config.search_mode )
+    if config.arch == 'resnet':
+      return SearchShapeImagenetResNet(config.block_name, config.layers, config.deep_stem, config.class_num)
+    else:
+      raise ValueError('invalid model config : {:}'.format(config))
+  else:
+    raise ValueError('invalid dataset in the model config : {:}'.format(config))
+
+
+def load_net_from_checkpoint(checkpoint):
+  assert osp.isfile(checkpoint), 'checkpoint {:} does not exist'.format(checkpoint)
+  checkpoint   = torch.load(checkpoint)
+  model_config = dict2config(checkpoint['model-config'], None)
+  model        = obtain_model(model_config)
+  model.load_state_dict(checkpoint['base-model'])
+  return model

lib/models/clone_weights.py

@@ -0,0 +1,62 @@
+import torch
+import torch.nn as nn
+
+
+def copy_conv(module, init):
+  assert isinstance(module, nn.Conv2d), 'invalid module : {:}'.format(module)
+  assert isinstance(init  , nn.Conv2d), 'invalid module : {:}'.format(init)
+  new_i, new_o = module.in_channels, module.out_channels
+  module.weight.copy_( init.weight.detach()[:new_o, :new_i] )
+  if module.bias is not None:
+    module.bias.copy_( init.bias.detach()[:new_o] )
+
+def copy_bn  (module, init):
+  assert isinstance(module, nn.BatchNorm2d), 'invalid module : {:}'.format(module)
+  assert isinstance(init  , nn.BatchNorm2d), 'invalid module : {:}'.format(init)
+  num_features = module.num_features
+  if module.weight is not None:
+    module.weight.copy_( init.weight.detach()[:num_features] )
+  if module.bias is not None:
+    module.bias.copy_( init.bias.detach()[:num_features] )
+  if module.running_mean is not None:
+    module.running_mean.copy_( init.running_mean.detach()[:num_features] )
+  if module.running_var  is not None:
+    module.running_var.copy_( init.running_var.detach()[:num_features] )
+
+def copy_fc  (module, init):
+  assert isinstance(module, nn.Linear), 'invalid module : {:}'.format(module)
+  assert isinstance(init  , nn.Linear), 'invalid module : {:}'.format(init)
+  new_i, new_o = module.in_features, module.out_features
+  module.weight.copy_( init.weight.detach()[:new_o, :new_i] )
+  if module.bias is not None:
+    module.bias.copy_( init.bias.detach()[:new_o] )
+
+def copy_base(module, init):
+  assert type(module).__name__ in ['ConvBNReLU', 'Downsample'], 'invalid module : {:}'.format(module)
+  assert type(  init).__name__ in ['ConvBNReLU', 'Downsample'], 'invalid module : {:}'.format(  init)
+  if module.conv is not None:
+    copy_conv(module.conv, init.conv)
+  if module.bn is not None:
+    copy_bn  (module.bn, init.bn)
+
+def copy_basic(module, init):
+  copy_base(module.conv_a, init.conv_a)
+  copy_base(module.conv_b, init.conv_b)
+  if module.downsample is not None:
+    if init.downsample is not None:
+      copy_base(module.downsample, init.downsample)
+    #else:
+    # import pdb; pdb.set_trace()
+
+
+def init_from_model(network, init_model):
+  with torch.no_grad():
+    copy_fc(network.classifier, init_model.classifier)
+    for base, target in zip(init_model.layers, network.layers):
+      assert type(base).__name__  == type(target).__name__, 'invalid type : {:} vs {:}'.format(base, target)
+      if type(base).__name__ == 'ConvBNReLU':
+        copy_base(target, base)
+      elif type(base).__name__ == 'ResNetBasicblock':
+        copy_basic(target, base)
+      else:
+        raise ValueError('unknown type name : {:}'.format( type(base).__name__ ))

lib/models/infers/InferCifarResNet.py

@@ -0,0 +1,166 @@
+import math, torch
+import torch.nn as nn
+import torch.nn.functional as F
+from ..initialization import initialize_resnet
+from ..SharedUtils    import additive_func
+
+
+class ConvBNReLU(nn.Module):
+  
+  def __init__(self, nIn, nOut, kernel, stride, padding, bias, has_avg, has_bn, has_relu):
+    super(ConvBNReLU, self).__init__()
+    if has_avg : self.avg = nn.AvgPool2d(kernel_size=2, stride=2, padding=0)
+    else       : self.avg = None
+    self.conv = nn.Conv2d(nIn, nOut, kernel_size=kernel, stride=stride, padding=padding, dilation=1, groups=1, bias=bias)
+    if has_bn  : self.bn  = nn.BatchNorm2d(nOut)
+    else       : self.bn  = None
+    if has_relu: self.relu = nn.ReLU(inplace=True)
+    else       : self.relu = None
+
+  def forward(self, inputs):
+    if self.avg : out = self.avg( inputs )
+    else        : out = inputs
+    conv = self.conv( out )
+    if self.bn  : out = self.bn( conv )
+    else        : out = conv
+    if self.relu: out = self.relu( out )
+    else        : out = out
+
+    return out
+
+
+class ResNetBasicblock(nn.Module):
+  num_conv  = 2
+  expansion = 1
+  def __init__(self, iCs, stride):
+    super(ResNetBasicblock, self).__init__()
+    assert stride == 1 or stride == 2, 'invalid stride {:}'.format(stride)
+    assert isinstance(iCs, tuple) or isinstance(iCs, list), 'invalid type of iCs : {:}'.format( iCs )
+    assert len(iCs) == 3,'invalid lengths of iCs : {:}'.format(iCs)
+    
+    self.conv_a = ConvBNReLU(iCs[0], iCs[1], 3, stride, 1, False, has_avg=False, has_bn=True, has_relu=True)
+    self.conv_b = ConvBNReLU(iCs[1], iCs[2], 3,      1, 1, False, has_avg=False, has_bn=True, has_relu=False)
+    residual_in = iCs[0]
+    if stride == 2:
+      self.downsample = ConvBNReLU(iCs[0], iCs[2], 1, 1, 0, False, has_avg=True, has_bn=False, has_relu=False)
+      residual_in = iCs[2]
+    elif iCs[0] != iCs[2]:
+      self.downsample = ConvBNReLU(iCs[0], iCs[2], 1, 1, 0, False, has_avg=False,has_bn=True , has_relu=False)
+    else:
+      self.downsample = None
+    #self.out_dim  = max(residual_in, iCs[2])
+    self.out_dim  = iCs[2]
+
+  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
+    out = residual + basicblock
+    return F.relu(out, inplace=True)
+
+
+
+class ResNetBottleneck(nn.Module):
+  expansion = 4
+  num_conv  = 3
+  def __init__(self, iCs, stride):
+    super(ResNetBottleneck, self).__init__()
+    assert stride == 1 or stride == 2, 'invalid stride {:}'.format(stride)
+    assert isinstance(iCs, tuple) or isinstance(iCs, list), 'invalid type of iCs : {:}'.format( iCs )
+    assert len(iCs) == 4,'invalid lengths of iCs : {:}'.format(iCs)
+    self.conv_1x1 = ConvBNReLU(iCs[0], iCs[1], 1,      1, 0, False, has_avg=False, has_bn=True, has_relu=True)
+    self.conv_3x3 = ConvBNReLU(iCs[1], iCs[2], 3, stride, 1, False, has_avg=False, has_bn=True, has_relu=True)
+    self.conv_1x4 = ConvBNReLU(iCs[2], iCs[3], 1,      1, 0, False, has_avg=False, has_bn=True, has_relu=False)
+    residual_in = iCs[0]
+    if stride == 2:
+      self.downsample = ConvBNReLU(iCs[0], iCs[3], 1, 1, 0, False, has_avg=True , has_bn=False, has_relu=False)
+      residual_in     = iCs[3]
+    elif iCs[0] != iCs[3]:
+      self.downsample = ConvBNReLU(iCs[0], iCs[3], 1, 1, 0, False, has_avg=False, has_bn=False, has_relu=False)
+      residual_in     = iCs[3]
+    else:
+      self.downsample = None
+    #self.out_dim = max(residual_in, iCs[3])
+    self.out_dim = iCs[3]
+
+  def forward(self, inputs):
+
+    bottleneck = self.conv_1x1(inputs)
+    bottleneck = self.conv_3x3(bottleneck)
+    bottleneck = self.conv_1x4(bottleneck)
+
+    if self.downsample is not None:
+      residual = self.downsample(inputs)
+    else:
+      residual = inputs
+    out = residual + bottleneck
+    return F.relu(out, inplace=True)
+
+
+
+class InferCifarResNet(nn.Module):
+
+  def __init__(self, block_name, depth, xblocks, xchannels, num_classes, zero_init_residual):
+    super(InferCifarResNet, self).__init__()
+
+    #Model type specifies number of layers for CIFAR-10 and CIFAR-100 model
+    if block_name == 'ResNetBasicblock':
+      block = ResNetBasicblock
+      assert (depth - 2) % 6 == 0, 'depth should be one of 20, 32, 44, 56, 110'
+      layer_blocks = (depth - 2) // 6
+    elif block_name == 'ResNetBottleneck':
+      block = ResNetBottleneck
+      assert (depth - 2) % 9 == 0, 'depth should be one of 164'
+      layer_blocks = (depth - 2) // 9
+    else:
+      raise ValueError('invalid block : {:}'.format(block_name))
+    assert len(xblocks) == 3, 'invalid xblocks : {:}'.format(xblocks)
+
+    self.message     = 'InferWidthCifarResNet : Depth : {:} , Layers for each block : {:}'.format(depth, layer_blocks)
+    self.num_classes = num_classes
+    self.xchannels   = xchannels
+    self.layers      = nn.ModuleList( [ ConvBNReLU(xchannels[0], xchannels[1], 3, 1, 1, False, has_avg=False, has_bn=True, has_relu=True) ] )
+    last_channel_idx = 1
+    for stage in range(3):
+      for iL in range(layer_blocks):
+        num_conv = block.num_conv 
+        iCs      = self.xchannels[last_channel_idx:last_channel_idx+num_conv+1]
+        stride   = 2 if stage > 0 and iL == 0 else 1
+        module   = block(iCs, stride)
+        last_channel_idx += num_conv
+        self.xchannels[last_channel_idx] = module.out_dim
+        self.layers.append  ( module )
+        self.message += "\nstage={:}, ilayer={:02d}/{:02d}, block={:03d}, iCs={:}, oC={:3d}, stride={:}".format(stage, iL, layer_blocks, len(self.layers)-1, iCs, module.out_dim, stride)
+        if iL + 1 == xblocks[stage]: # reach the maximum depth
+          out_channel = module.out_dim
+          for iiL in range(iL+1, layer_blocks):
+            last_channel_idx += num_conv
+          self.xchannels[last_channel_idx] = module.out_dim
+          break
+  
+    self.avgpool    = nn.AvgPool2d(8)
+    self.classifier = nn.Linear(self.xchannels[-1], num_classes)
+    
+    self.apply(initialize_resnet)
+    if zero_init_residual:
+      for m in self.modules():
+        if isinstance(m, ResNetBasicblock):
+          nn.init.constant_(m.conv_b.bn.weight, 0)
+        elif isinstance(m, ResNetBottleneck):
+          nn.init.constant_(m.conv_1x4.bn.weight, 0)
+
+  def get_message(self):
+    return self.message
+
+  def forward(self, inputs):
+    x = inputs
+    for i, layer in enumerate(self.layers):
+      x = layer( x )
+    features = self.avgpool(x)
+    features = features.view(features.size(0), -1)
+    logits   = self.classifier(features)
+    return features, logits

lib/models/infers/InferCifarResNet_depth.py

@@ -0,0 +1,149 @@
+import math, torch
+import torch.nn as nn
+import torch.nn.functional as F
+from ..initialization import initialize_resnet
+from ..SharedUtils    import additive_func
+
+
+class ConvBNReLU(nn.Module):
+  
+  def __init__(self, nIn, nOut, kernel, stride, padding, bias, has_avg, has_bn, has_relu):
+    super(ConvBNReLU, self).__init__()
+    if has_avg : self.avg = nn.AvgPool2d(kernel_size=2, stride=2, padding=0)
+    else       : self.avg = None
+    self.conv = nn.Conv2d(nIn, nOut, kernel_size=kernel, stride=stride, padding=padding, dilation=1, groups=1, bias=bias)
+    if has_bn  : self.bn  = nn.BatchNorm2d(nOut)
+    else       : self.bn  = None
+    if has_relu: self.relu = nn.ReLU(inplace=True)
+    else       : self.relu = None
+
+  def forward(self, inputs):
+    if self.avg : out = self.avg( inputs )
+    else        : out = inputs
+    conv = self.conv( out )
+    if self.bn  : out = self.bn( conv )
+    else        : out = conv
+    if self.relu: out = self.relu( out )
+    else        : out = out
+
+    return out
+
+
+class ResNetBasicblock(nn.Module):
+  num_conv  = 2
+  expansion = 1
+  def __init__(self, inplanes, planes, stride):
+    super(ResNetBasicblock, self).__init__()
+    assert stride == 1 or stride == 2, 'invalid stride {:}'.format(stride)
+    
+    self.conv_a = ConvBNReLU(inplanes, planes, 3, stride, 1, False, has_avg=False, has_bn=True, has_relu=True)
+    self.conv_b = ConvBNReLU(  planes, planes, 3,      1, 1, False, has_avg=False, has_bn=True, has_relu=False)
+    if stride == 2:
+      self.downsample = ConvBNReLU(inplanes, planes, 1, 1, 0, False, has_avg=True, has_bn=False, has_relu=False)
+    elif inplanes != planes:
+      self.downsample = ConvBNReLU(inplanes, planes, 1, 1, 0, False, has_avg=False,has_bn=True , has_relu=False)
+    else:
+      self.downsample = None
+    self.out_dim  = planes
+
+  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
+    out = residual + basicblock
+    return F.relu(out, inplace=True)
+
+
+
+class ResNetBottleneck(nn.Module):
+  expansion = 4
+  num_conv  = 3
+  def __init__(self, inplanes, planes, stride):
+    super(ResNetBottleneck, self).__init__()
+    assert stride == 1 or stride == 2, 'invalid stride {:}'.format(stride)
+    self.conv_1x1 = ConvBNReLU(inplanes, planes, 1,      1, 0, False, has_avg=False, has_bn=True, has_relu=True)
+    self.conv_3x3 = ConvBNReLU(  planes, planes, 3, stride, 1, False, has_avg=False, has_bn=True, has_relu=True)
+    self.conv_1x4 = ConvBNReLU(planes, planes*self.expansion, 1,      1, 0, False, has_avg=False, has_bn=True, has_relu=False)
+    if stride == 2:
+      self.downsample = ConvBNReLU(inplanes, planes*self.expansion, 1, 1, 0, False, has_avg=True , has_bn=False, has_relu=False)
+    elif inplanes != planes*self.expansion:
+      self.downsample = ConvBNReLU(inplanes, planes*self.expansion, 1, 1, 0, False, has_avg=False, has_bn=False, has_relu=False)
+    else:
+      self.downsample = None
+    self.out_dim = planes*self.expansion
+
+  def forward(self, inputs):
+
+    bottleneck = self.conv_1x1(inputs)
+    bottleneck = self.conv_3x3(bottleneck)
+    bottleneck = self.conv_1x4(bottleneck)
+
+    if self.downsample is not None:
+      residual = self.downsample(inputs)
+    else:
+      residual = inputs
+    out = residual + bottleneck
+    return F.relu(out, inplace=True)
+
+
+
+class InferDepthCifarResNet(nn.Module):
+
+  def __init__(self, block_name, depth, xblocks, num_classes, zero_init_residual):
+    super(InferDepthCifarResNet, self).__init__()
+
+    #Model type specifies number of layers for CIFAR-10 and CIFAR-100 model
+    if block_name == 'ResNetBasicblock':
+      block = ResNetBasicblock
+      assert (depth - 2) % 6 == 0, 'depth should be one of 20, 32, 44, 56, 110'
+      layer_blocks = (depth - 2) // 6
+    elif block_name == 'ResNetBottleneck':
+      block = ResNetBottleneck
+      assert (depth - 2) % 9 == 0, 'depth should be one of 164'
+      layer_blocks = (depth - 2) // 9
+    else:
+      raise ValueError('invalid block : {:}'.format(block_name))
+    assert len(xblocks) == 3, 'invalid xblocks : {:}'.format(xblocks)
+
+    self.message     = 'InferWidthCifarResNet : Depth : {:} , Layers for each block : {:}'.format(depth, layer_blocks)
+    self.num_classes = num_classes
+    self.layers      = nn.ModuleList( [ ConvBNReLU(3, 16, 3, 1, 1, False, has_avg=False, has_bn=True, has_relu=True) ] )
+    self.channels    = [16]
+    for stage in range(3):
+      for iL in range(layer_blocks):
+        iC       = self.channels[-1]
+        planes = 16 * (2**stage)
+        stride   = 2 if stage > 0 and iL == 0 else 1
+        module   = block(iC, planes, stride)
+        self.channels.append( module.out_dim )
+        self.layers.append  ( module )
+        self.message += "\nstage={:}, ilayer={:02d}/{:02d}, block={:03d}, iC={:}, oC={:3d}, stride={:}".format(stage, iL, layer_blocks, len(self.layers)-1, planes, module.out_dim, stride)
+        if iL + 1 == xblocks[stage]: # reach the maximum depth
+          break
+  
+    self.avgpool    = nn.AvgPool2d(8)
+    self.classifier = nn.Linear(self.channels[-1], num_classes)
+    
+    self.apply(initialize_resnet)
+    if zero_init_residual:
+      for m in self.modules():
+        if isinstance(m, ResNetBasicblock):
+          nn.init.constant_(m.conv_b.bn.weight, 0)
+        elif isinstance(m, ResNetBottleneck):
+          nn.init.constant_(m.conv_1x4.bn.weight, 0)
+
+  def get_message(self):
+    return self.message
+
+  def forward(self, inputs):
+    x = inputs
+    for i, layer in enumerate(self.layers):
+      x = layer( x )
+    features = self.avgpool(x)
+    features = features.view(features.size(0), -1)
+    logits   = self.classifier(features)
+    return features, logits

lib/models/infers/InferCifarResNet_width.py

@@ -0,0 +1,159 @@
+import math, torch
+import torch.nn as nn
+import torch.nn.functional as F
+from ..initialization import initialize_resnet
+from ..SharedUtils    import additive_func
+
+
+class ConvBNReLU(nn.Module):
+  
+  def __init__(self, nIn, nOut, kernel, stride, padding, bias, has_avg, has_bn, has_relu):
+    super(ConvBNReLU, self).__init__()
+    if has_avg : self.avg = nn.AvgPool2d(kernel_size=2, stride=2, padding=0)
+    else       : self.avg = None
+    self.conv = nn.Conv2d(nIn, nOut, kernel_size=kernel, stride=stride, padding=padding, dilation=1, groups=1, bias=bias)
+    if has_bn  : self.bn  = nn.BatchNorm2d(nOut)
+    else       : self.bn  = None
+    if has_relu: self.relu = nn.ReLU(inplace=True)
+    else       : self.relu = None
+
+  def forward(self, inputs):
+    if self.avg : out = self.avg( inputs )
+    else        : out = inputs
+    conv = self.conv( out )
+    if self.bn  : out = self.bn( conv )
+    else        : out = conv
+    if self.relu: out = self.relu( out )
+    else        : out = out
+
+    return out
+
+
+class ResNetBasicblock(nn.Module):
+  num_conv  = 2
+  expansion = 1
+  def __init__(self, iCs, stride):
+    super(ResNetBasicblock, self).__init__()
+    assert stride == 1 or stride == 2, 'invalid stride {:}'.format(stride)
+    assert isinstance(iCs, tuple) or isinstance(iCs, list), 'invalid type of iCs : {:}'.format( iCs )
+    assert len(iCs) == 3,'invalid lengths of iCs : {:}'.format(iCs)
+    
+    self.conv_a = ConvBNReLU(iCs[0], iCs[1], 3, stride, 1, False, has_avg=False, has_bn=True, has_relu=True)
+    self.conv_b = ConvBNReLU(iCs[1], iCs[2], 3,      1, 1, False, has_avg=False, has_bn=True, has_relu=False)
+    residual_in = iCs[0]
+    if stride == 2:
+      self.downsample = ConvBNReLU(iCs[0], iCs[2], 1, 1, 0, False, has_avg=True, has_bn=False, has_relu=False)
+      residual_in = iCs[2]
+    elif iCs[0] != iCs[2]:
+      self.downsample = ConvBNReLU(iCs[0], iCs[2], 1, 1, 0, False, has_avg=False,has_bn=True , has_relu=False)
+    else:
+      self.downsample = None
+    #self.out_dim  = max(residual_in, iCs[2])
+    self.out_dim  = iCs[2]
+
+  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
+    out = residual + basicblock
+    return F.relu(out, inplace=True)
+
+
+
+class ResNetBottleneck(nn.Module):
+  expansion = 4
+  num_conv  = 3
+  def __init__(self, iCs, stride):
+    super(ResNetBottleneck, self).__init__()
+    assert stride == 1 or stride == 2, 'invalid stride {:}'.format(stride)
+    assert isinstance(iCs, tuple) or isinstance(iCs, list), 'invalid type of iCs : {:}'.format( iCs )
+    assert len(iCs) == 4,'invalid lengths of iCs : {:}'.format(iCs)
+    self.conv_1x1 = ConvBNReLU(iCs[0], iCs[1], 1,      1, 0, False, has_avg=False, has_bn=True, has_relu=True)
+    self.conv_3x3 = ConvBNReLU(iCs[1], iCs[2], 3, stride, 1, False, has_avg=False, has_bn=True, has_relu=True)
+    self.conv_1x4 = ConvBNReLU(iCs[2], iCs[3], 1,      1, 0, False, has_avg=False, has_bn=True, has_relu=False)
+    residual_in = iCs[0]
+    if stride == 2:
+      self.downsample = ConvBNReLU(iCs[0], iCs[3], 1, 1, 0, False, has_avg=True , has_bn=False, has_relu=False)
+      residual_in     = iCs[3]
+    elif iCs[0] != iCs[3]:
+      self.downsample = ConvBNReLU(iCs[0], iCs[3], 1, 1, 0, False, has_avg=False, has_bn=False, has_relu=False)
+      residual_in     = iCs[3]
+    else:
+      self.downsample = None
+    #self.out_dim = max(residual_in, iCs[3])
+    self.out_dim = iCs[3]
+
+  def forward(self, inputs):
+
+    bottleneck = self.conv_1x1(inputs)
+    bottleneck = self.conv_3x3(bottleneck)
+    bottleneck = self.conv_1x4(bottleneck)
+
+    if self.downsample is not None:
+      residual = self.downsample(inputs)
+    else:
+      residual = inputs
+    out = residual + bottleneck
+    return F.relu(out, inplace=True)
+
+
+
+class InferWidthCifarResNet(nn.Module):
+
+  def __init__(self, block_name, depth, xchannels, num_classes, zero_init_residual):
+    super(InferWidthCifarResNet, self).__init__()
+
+    #Model type specifies number of layers for CIFAR-10 and CIFAR-100 model
+    if block_name == 'ResNetBasicblock':
+      block = ResNetBasicblock
+      assert (depth - 2) % 6 == 0, 'depth should be one of 20, 32, 44, 56, 110'
+      layer_blocks = (depth - 2) // 6
+    elif block_name == 'ResNetBottleneck':
+      block = ResNetBottleneck
+      assert (depth - 2) % 9 == 0, 'depth should be one of 164'
+      layer_blocks = (depth - 2) // 9
+    else:
+      raise ValueError('invalid block : {:}'.format(block_name))
+
+    self.message     = 'InferWidthCifarResNet : Depth : {:} , Layers for each block : {:}'.format(depth, layer_blocks)
+    self.num_classes = num_classes
+    self.xchannels   = xchannels
+    self.layers      = nn.ModuleList( [ ConvBNReLU(xchannels[0], xchannels[1], 3, 1, 1, False, has_avg=False, has_bn=True, has_relu=True) ] )
+    last_channel_idx = 1
+    for stage in range(3):
+      for iL in range(layer_blocks):
+        num_conv = block.num_conv 
+        iCs      = self.xchannels[last_channel_idx:last_channel_idx+num_conv+1]
+        stride   = 2 if stage > 0 and iL == 0 else 1
+        module   = block(iCs, stride)
+        last_channel_idx += num_conv
+        self.xchannels[last_channel_idx] = module.out_dim
+        self.layers.append  ( module )
+        self.message += "\nstage={:}, ilayer={:02d}/{:02d}, block={:03d}, iCs={:}, oC={:3d}, stride={:}".format(stage, iL, layer_blocks, len(self.layers)-1, iCs, module.out_dim, stride)
+  
+    self.avgpool    = nn.AvgPool2d(8)
+    self.classifier = nn.Linear(self.xchannels[-1], num_classes)
+    
+    self.apply(initialize_resnet)
+    if zero_init_residual:
+      for m in self.modules():
+        if isinstance(m, ResNetBasicblock):
+          nn.init.constant_(m.conv_b.bn.weight, 0)
+        elif isinstance(m, ResNetBottleneck):
+          nn.init.constant_(m.conv_1x4.bn.weight, 0)
+
+  def get_message(self):
+    return self.message
+
+  def forward(self, inputs):
+    x = inputs
+    for i, layer in enumerate(self.layers):
+      x = layer( x )
+    features = self.avgpool(x)
+    features = features.view(features.size(0), -1)
+    logits   = self.classifier(features)
+    return features, logits

lib/models/infers/InferImagenetResNet.py

@@ -0,0 +1,169 @@
+import math, torch
+import torch.nn as nn
+import torch.nn.functional as F
+from ..initialization import initialize_resnet
+from ..SharedUtils    import additive_func
+
+
+class ConvBNReLU(nn.Module):
+  
+  num_conv  = 1
+  def __init__(self, nIn, nOut, kernel, stride, padding, bias, has_avg, has_bn, has_relu):
+    super(ConvBNReLU, self).__init__()
+    if has_avg : self.avg = nn.AvgPool2d(kernel_size=2, stride=2, padding=0)
+    else       : self.avg = None
+    self.conv = nn.Conv2d(nIn, nOut, kernel_size=kernel, stride=stride, padding=padding, dilation=1, groups=1, bias=bias)
+    if has_bn  : self.bn  = nn.BatchNorm2d(nOut)
+    else       : self.bn  = None
+    if has_relu: self.relu = nn.ReLU(inplace=True)
+    else       : self.relu = None
+
+  def forward(self, inputs):
+    if self.avg : out = self.avg( inputs )
+    else        : out = inputs
+    conv = self.conv( out )
+    if self.bn  : out = self.bn( conv )
+    else        : out = conv
+    if self.relu: out = self.relu( out )
+    else        : out = out
+
+    return out
+
+
+class ResNetBasicblock(nn.Module):
+  num_conv  = 2
+  expansion = 1
+  def __init__(self, iCs, stride):
+    super(ResNetBasicblock, self).__init__()
+    assert stride == 1 or stride == 2, 'invalid stride {:}'.format(stride)
+    assert isinstance(iCs, tuple) or isinstance(iCs, list), 'invalid type of iCs : {:}'.format( iCs )
+    assert len(iCs) == 3,'invalid lengths of iCs : {:}'.format(iCs)
+    
+    self.conv_a = ConvBNReLU(iCs[0], iCs[1], 3, stride, 1, False, has_avg=False, has_bn=True, has_relu=True)
+    self.conv_b = ConvBNReLU(iCs[1], iCs[2], 3,      1, 1, False, has_avg=False, has_bn=True, has_relu=False)
+    residual_in = iCs[0]
+    if stride == 2:
+      self.downsample = ConvBNReLU(iCs[0], iCs[2], 1, 1, 0, False, has_avg=True, has_bn=True, has_relu=False)
+      residual_in = iCs[2]
+    elif iCs[0] != iCs[2]:
+      self.downsample = ConvBNReLU(iCs[0], iCs[2], 1, 1, 0, False, has_avg=False,has_bn=True , has_relu=False)
+    else:
+      self.downsample = None
+    #self.out_dim  = max(residual_in, iCs[2])
+    self.out_dim  = iCs[2]
+
+  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
+    out = residual + basicblock
+    return F.relu(out, inplace=True)
+
+
+
+class ResNetBottleneck(nn.Module):
+  expansion = 4
+  num_conv  = 3
+  def __init__(self, iCs, stride):
+    super(ResNetBottleneck, self).__init__()
+    assert stride == 1 or stride == 2, 'invalid stride {:}'.format(stride)
+    assert isinstance(iCs, tuple) or isinstance(iCs, list), 'invalid type of iCs : {:}'.format( iCs )
+    assert len(iCs) == 4,'invalid lengths of iCs : {:}'.format(iCs)
+    self.conv_1x1 = ConvBNReLU(iCs[0], iCs[1], 1,      1, 0, False, has_avg=False, has_bn=True, has_relu=True)
+    self.conv_3x3 = ConvBNReLU(iCs[1], iCs[2], 3, stride, 1, False, has_avg=False, has_bn=True, has_relu=True)
+    self.conv_1x4 = ConvBNReLU(iCs[2], iCs[3], 1,      1, 0, False, has_avg=False, has_bn=True, has_relu=False)
+    residual_in = iCs[0]
+    if stride == 2:
+      self.downsample = ConvBNReLU(iCs[0], iCs[3], 1, 1, 0, False, has_avg=True , has_bn=True, has_relu=False)
+      residual_in     = iCs[3]
+    elif iCs[0] != iCs[3]:
+      self.downsample = ConvBNReLU(iCs[0], iCs[3], 1, 1, 0, False, has_avg=False, has_bn=True, has_relu=False)
+      residual_in     = iCs[3]
+    else:
+      self.downsample = None
+    #self.out_dim = max(residual_in, iCs[3])
+    self.out_dim = iCs[3]
+
+  def forward(self, inputs):
+
+    bottleneck = self.conv_1x1(inputs)
+    bottleneck = self.conv_3x3(bottleneck)
+    bottleneck = self.conv_1x4(bottleneck)
+
+    if self.downsample is not None:
+      residual = self.downsample(inputs)
+    else:
+      residual = inputs
+    out = residual + bottleneck
+    return F.relu(out, inplace=True)
+
+
+
+class InferImagenetResNet(nn.Module):
+
+  def __init__(self, block_name, layers, xblocks, xchannels, deep_stem, num_classes, zero_init_residual):
+    super(InferImagenetResNet, self).__init__()
+
+    #Model type specifies number of layers for CIFAR-10 and CIFAR-100 model
+    if block_name == 'BasicBlock':
+      block = ResNetBasicblock
+    elif block_name == 'Bottleneck':
+      block = ResNetBottleneck
+    else:
+      raise ValueError('invalid block : {:}'.format(block_name))
+    assert len(xblocks) == len(layers), 'invalid layers : {:} vs xblocks : {:}'.format(layers, xblocks)
+
+    self.message     = 'InferImagenetResNet : Depth : {:} -> {:}, Layers for each block : {:}'.format(sum(layers)*block.num_conv, sum(xblocks)*block.num_conv, xblocks)
+    self.num_classes = num_classes
+    self.xchannels   = xchannels
+    if not deep_stem:
+      self.layers      = nn.ModuleList( [ ConvBNReLU(xchannels[0], xchannels[1], 7, 2, 3, False, has_avg=False, has_bn=True, has_relu=True) ] )
+      last_channel_idx = 1
+    else:
+      self.layers      = nn.ModuleList( [ ConvBNReLU(xchannels[0], xchannels[1], 3, 2, 1, False, has_avg=False, has_bn=True, has_relu=True)
+                                         ,ConvBNReLU(xchannels[1], xchannels[2], 3, 1, 1, False, has_avg=False, has_bn=True, has_relu=True) ] )
+      last_channel_idx = 2
+    self.layers.append( nn.MaxPool2d(kernel_size=3, stride=2, padding=1) )
+    for stage, layer_blocks in enumerate(layers):
+      for iL in range(layer_blocks):
+        num_conv = block.num_conv 
+        iCs      = self.xchannels[last_channel_idx:last_channel_idx+num_conv+1]
+        stride   = 2 if stage > 0 and iL == 0 else 1
+        module   = block(iCs, stride)
+        last_channel_idx += num_conv
+        self.xchannels[last_channel_idx] = module.out_dim
+        self.layers.append  ( module )
+        self.message += "\nstage={:}, ilayer={:02d}/{:02d}, block={:03d}, iCs={:}, oC={:3d}, stride={:}".format(stage, iL, layer_blocks, len(self.layers)-1, iCs, module.out_dim, stride)
+        if iL + 1 == xblocks[stage]: # reach the maximum depth
+          out_channel = module.out_dim
+          for iiL in range(iL+1, layer_blocks):
+            last_channel_idx += num_conv
+          self.xchannels[last_channel_idx] = module.out_dim
+          break
+    assert last_channel_idx + 1 == len(self.xchannels), '{:} vs {:}'.format(last_channel_idx, len(self.xchannels))
+    self.avgpool    = nn.AdaptiveAvgPool2d((1,1))
+    self.classifier = nn.Linear(self.xchannels[-1], num_classes)
+    
+    self.apply(initialize_resnet)
+    if zero_init_residual:
+      for m in self.modules():
+        if isinstance(m, ResNetBasicblock):
+          nn.init.constant_(m.conv_b.bn.weight, 0)
+        elif isinstance(m, ResNetBottleneck):
+          nn.init.constant_(m.conv_1x4.bn.weight, 0)
+
+  def get_message(self):
+    return self.message
+
+  def forward(self, inputs):
+    x = inputs
+    for i, layer in enumerate(self.layers):
+      x = layer( x )
+    features = self.avgpool(x)
+    features = features.view(features.size(0), -1)
+    logits   = self.classifier(features)
+    return features, logits

lib/models/infers/InferMobileNetV2.py

@@ -0,0 +1,119 @@
+# MobileNetV2: Inverted Residuals and Linear Bottlenecks, CVPR 2018
+from torch import nn
+from ..initialization import initialize_resnet
+from ..SharedUtils    import additive_func, parse_channel_info
+
+
+class ConvBNReLU(nn.Module):
+  def __init__(self, in_planes, out_planes, kernel_size, stride, groups, has_bn=True, has_relu=True):
+    super(ConvBNReLU, self).__init__()
+    padding = (kernel_size - 1) // 2
+    self.conv = nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False)
+    if has_bn: self.bn = nn.BatchNorm2d(out_planes)
+    else     : self.bn = None
+    if has_relu: self.relu = nn.ReLU6(inplace=True)
+    else       : self.relu = None
+  
+  def forward(self, x):
+    out = self.conv( x )
+    if self.bn:   out = self.bn  ( out )
+    if self.relu: out = self.relu( out )
+    return out
+
+
+class InvertedResidual(nn.Module):
+  def __init__(self, channels, stride, expand_ratio, additive):
+    super(InvertedResidual, self).__init__()
+    self.stride = stride
+    assert stride in [1, 2], 'invalid stride : {:}'.format(stride)
+    assert len(channels) in [2, 3], 'invalid channels : {:}'.format(channels)
+
+    if len(channels) == 2:
+      layers = []
+    else:
+      layers = [ConvBNReLU(channels[0], channels[1], 1, 1, 1)]
+    layers.extend([
+      # dw
+      ConvBNReLU(channels[-2], channels[-2], 3, stride, channels[-2]),
+      # pw-linear
+      ConvBNReLU(channels[-2], channels[-1], 1, 1, 1, True, False),
+    ])
+    self.conv = nn.Sequential(*layers)
+    self.additive = additive
+    if self.additive and channels[0] != channels[-1]:
+      self.shortcut = ConvBNReLU(channels[0], channels[-1], 1, 1, 1, True, False)
+    else:
+      self.shortcut = None
+    self.out_dim  = channels[-1]
+
+  def forward(self, x):
+    out = self.conv(x)
+    # if self.additive: return additive_func(out, x)
+    if self.shortcut: return out + self.shortcut(x)
+    else            : return out
+
+
+class InferMobileNetV2(nn.Module):
+  def __init__(self, num_classes, xchannels, xblocks, dropout):
+    super(InferMobileNetV2, self).__init__()
+    block = InvertedResidual
+    inverted_residual_setting = [
+      # t, c,  n, s
+      [1, 16 , 1, 1],
+      [6, 24 , 2, 2],
+      [6, 32 , 3, 2],
+      [6, 64 , 4, 2],
+      [6, 96 , 3, 1],
+      [6, 160, 3, 2],
+      [6, 320, 1, 1],
+    ]
+    assert len(inverted_residual_setting) == len(xblocks), 'invalid number of layers : {:} vs {:}'.format(len(inverted_residual_setting), len(xblocks))
+    for block_num, ir_setting in zip(xblocks, inverted_residual_setting):
+      assert block_num <= ir_setting[2], '{:} vs {:}'.format(block_num, ir_setting)
+    xchannels = parse_channel_info(xchannels)
+    #for i, chs in enumerate(xchannels):
+    #  if i > 0: assert chs[0] == xchannels[i-1][-1], 'Layer[{:}] is invalid {:} vs {:}'.format(i, xchannels[i-1], chs)
+    self.xchannels = xchannels
+    self.message     = 'InferMobileNetV2 : xblocks={:}'.format(xblocks)
+    # building first layer
+    features = [ConvBNReLU(xchannels[0][0], xchannels[0][1], 3, 2, 1)]
+    last_channel_idx = 1
+
+    # building inverted residual blocks
+    for stage, (t, c, n, s) in enumerate(inverted_residual_setting):
+      for i in range(n):
+        stride = s if i == 0 else 1
+        additv = True if i > 0 else False
+        module = block(self.xchannels[last_channel_idx], stride, t, additv)
+        features.append(module)
+        self.message += "\nstage={:}, ilayer={:02d}/{:02d}, block={:03d}, Cs={:}, stride={:}, expand={:}, original-C={:}".format(stage, i, n, len(features), self.xchannels[last_channel_idx], stride, t, c)
+        last_channel_idx += 1
+        if i + 1 == xblocks[stage]:
+          out_channel = module.out_dim
+          for iiL in range(i+1, n):
+            last_channel_idx += 1
+          self.xchannels[last_channel_idx][0] = module.out_dim
+          break
+    # building last several layers
+    features.append(ConvBNReLU(self.xchannels[last_channel_idx][0], self.xchannels[last_channel_idx][1], 1, 1, 1))
+    assert last_channel_idx + 2 == len(self.xchannels), '{:} vs {:}'.format(last_channel_idx, len(self.xchannels))
+    # make it nn.Sequential
+    self.features = nn.Sequential(*features)
+
+    # building classifier
+    self.classifier = nn.Sequential(
+      nn.Dropout(dropout),
+      nn.Linear(self.xchannels[last_channel_idx][1], num_classes),
+    )
+
+    # weight initialization
+    self.apply( initialize_resnet )
+
+  def get_message(self):
+    return self.message
+
+  def forward(self, inputs):
+    features = self.features(inputs)
+    vectors  = features.mean([2, 3])
+    predicts = self.classifier(vectors)
+    return features, predicts

lib/models/infers/__init__.py

@@ -0,0 +1,5 @@
+from .InferCifarResNet_width import InferWidthCifarResNet
+from .InferImagenetResNet    import InferImagenetResNet
+from .InferCifarResNet_depth import InferDepthCifarResNet
+from .InferCifarResNet       import InferCifarResNet
+from .InferMobileNetV2       import InferMobileNetV2

lib/models/infers/shared_utils.py

@@ -0,0 +1,7 @@
+# Xuanyi Dong
+
+def parse_channel_info(xstring):
+  blocks = xstring.split(' ')
+  blocks = [x.split('-') for x in blocks]
+  blocks = [[int(_) for _ in x] for x in blocks]
+  return blocks

lib/models/initialization.py

@@ -0,0 +1,18 @@
+import torch
+import torch.nn as nn
+
+
+def initialize_resnet(m):
+  if isinstance(m, nn.Conv2d):
+    nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+    if m.bias is not None:
+      nn.init.constant_(m.bias, 0)
+  elif isinstance(m, nn.BatchNorm2d):
+    nn.init.constant_(m.weight, 1)
+    if m.bias is not None:
+      nn.init.constant_(m.bias, 0)
+  elif isinstance(m, nn.Linear):
+    nn.init.normal_(m.weight, 0, 0.01)
+    nn.init.constant_(m.bias, 0)
+
+

lib/models/searchs/SearchCifarResNet.py

@@ -0,0 +1,502 @@
+import math, torch
+from collections import OrderedDict
+from bisect import bisect_right
+import torch.nn as nn
+from ..initialization import initialize_resnet
+from ..SharedUtils    import additive_func
+from .SoftSelect      import select2withP, ChannelWiseInter
+from .SoftSelect      import linear_forward
+from .SoftSelect      import get_width_choices
+
+
+def get_depth_choices(nDepth, return_num):
+  if nDepth == 2:
+    choices = (1, 2)
+  elif nDepth == 3:
+    choices = (1, 2, 3)
+  elif nDepth > 3:
+    choices = list(range(1, nDepth+1, 2))
+    if choices[-1] < nDepth: choices.append(nDepth)
+  else:
+    raise ValueError('invalid nDepth : {:}'.format(nDepth))
+  if return_num: return len(choices)
+  else         : return choices
+  
+
+
+
+def conv_forward(inputs, conv, choices):
+  iC = conv.in_channels
+  fill_size = list(inputs.size())
+  fill_size[1] = iC - fill_size[1]
+  filled  = torch.zeros(fill_size, device=inputs.device)
+  xinputs = torch.cat((inputs, filled), dim=1)
+  outputs = conv(xinputs)
+  selecteds = [outputs[:,:oC] for oC in choices]
+  return selecteds
+
+
+class ConvBNReLU(nn.Module):
+  num_conv  = 1
+  def __init__(self, nIn, nOut, kernel, stride, padding, bias, has_avg, has_bn, has_relu):
+    super(ConvBNReLU, self).__init__()
+    self.InShape  = None
+    self.OutShape = None
+    self.choices  = get_width_choices(nOut)
+    self.register_buffer('choices_tensor', torch.Tensor( self.choices ))
+
+    if has_avg : self.avg = nn.AvgPool2d(kernel_size=2, stride=2, padding=0)
+    else       : self.avg = None
+    self.conv = nn.Conv2d(nIn, nOut, kernel_size=kernel, stride=stride, padding=padding, dilation=1, groups=1, bias=bias)
+    #if has_bn  : self.bn  = nn.BatchNorm2d(nOut)
+    #else       : self.bn  = None
+    self.has_bn = has_bn
+    self.BNs  = nn.ModuleList()
+    for i, _out in enumerate(self.choices):
+      self.BNs.append(nn.BatchNorm2d(_out))
+    if has_relu: self.relu = nn.ReLU(inplace=True)
+    else       : self.relu = None
+    self.in_dim   = nIn
+    self.out_dim  = nOut
+    self.search_mode = 'basic'
+
+  def get_flops(self, channels, check_range=True, divide=1):
+    iC, oC = channels
+    if check_range: assert iC <= self.conv.in_channels and oC <= self.conv.out_channels, '{:} vs {:}  |  {:} vs {:}'.format(iC, self.conv.in_channels, oC, self.conv.out_channels)
+    assert isinstance(self.InShape, tuple) and len(self.InShape) == 2, 'invalid in-shape : {:}'.format(self.InShape)
+    assert isinstance(self.OutShape, tuple) and len(self.OutShape) == 2, 'invalid out-shape : {:}'.format(self.OutShape)
+    #conv_per_position_flops = self.conv.kernel_size[0] * self.conv.kernel_size[1] * iC * oC / self.conv.groups
+    conv_per_position_flops = (self.conv.kernel_size[0] * self.conv.kernel_size[1] * 1.0 / self.conv.groups)
+    all_positions = self.OutShape[0] * self.OutShape[1]
+    flops = (conv_per_position_flops * all_positions / divide) * iC * oC
+    if self.conv.bias is not None: flops += all_positions / divide
+    return flops
+
+  def get_range(self):
+    return [self.choices]
+
+  def forward(self, inputs):
+    if self.search_mode == 'basic':
+      return self.basic_forward(inputs)
+    elif self.search_mode == 'search':
+      return self.search_forward(inputs)
+    else:
+      raise ValueError('invalid search_mode = {:}'.format(self.search_mode))
+
+  def search_forward(self, tuple_inputs):
+    assert isinstance(tuple_inputs, tuple) and len(tuple_inputs) == 5, 'invalid type input : {:}'.format( type(tuple_inputs) )
+    inputs, expected_inC, probability, index, prob = tuple_inputs
+    index, prob = torch.squeeze(index).tolist(), torch.squeeze(prob)
+    probability = torch.squeeze(probability)
+    assert len(index) == 2, 'invalid length : {:}'.format(index)
+    # compute expected flop
+    #coordinates   = torch.arange(self.x_range[0], self.x_range[1]+1).type_as(probability)
+    expected_outC = (self.choices_tensor * probability).sum()
+    expected_flop = self.get_flops([expected_inC, expected_outC], False, 1e6)
+    if self.avg : out = self.avg( inputs )
+    else        : out = inputs
+    # convolutional layer
+    out_convs = conv_forward(out, self.conv, [self.choices[i] for i in index])
+    out_bns   = [self.BNs[idx](out_conv) for idx, out_conv in zip(index, out_convs)]
+    # merge
+    out_channel = max([x.size(1) for x in out_bns])
+    outA = ChannelWiseInter(out_bns[0], out_channel)
+    outB = ChannelWiseInter(out_bns[1], out_channel)
+    out  = outA * prob[0] + outB * prob[1]
+    #out = additive_func(out_bns[0]*prob[0], out_bns[1]*prob[1])
+
+    if self.relu: out = self.relu( out )
+    else        : out = out
+    return out, expected_outC, expected_flop
+
+  def basic_forward(self, inputs):
+    if self.avg : out = self.avg( inputs )
+    else        : out = inputs
+    conv = self.conv( out )
+    if self.has_bn:out= self.BNs[-1]( conv )
+    else        : out = conv
+    if self.relu: out = self.relu( out )
+    else        : out = out
+    if self.InShape is None:
+      self.InShape  = (inputs.size(-2), inputs.size(-1))
+      self.OutShape = (out.size(-2)   , out.size(-1))
+    return out
+
+
+class ResNetBasicblock(nn.Module):
+  expansion = 1
+  num_conv  = 2
+  def __init__(self, inplanes, planes, stride):
+    super(ResNetBasicblock, self).__init__()
+    assert stride == 1 or stride == 2, 'invalid stride {:}'.format(stride)
+    self.conv_a = ConvBNReLU(inplanes, planes, 3, stride, 1, False, has_avg=False, has_bn=True, has_relu=True)
+    self.conv_b = ConvBNReLU(  planes, planes, 3,      1, 1, False, has_avg=False, has_bn=True, has_relu=False)
+    if stride == 2:
+      self.downsample = ConvBNReLU(inplanes, planes, 1, 1, 0, False, has_avg=True, has_bn=False, has_relu=False)
+    elif inplanes != planes:
+      self.downsample = ConvBNReLU(inplanes, planes, 1, 1, 0, False, has_avg=False,has_bn=True , has_relu=False)
+    else:
+      self.downsample = None
+    self.out_dim     = planes
+    self.search_mode = 'basic'
+
+  def get_range(self):
+    return self.conv_a.get_range() + self.conv_b.get_range()
+
+  def get_flops(self, channels):
+    assert len(channels) == 3, 'invalid channels : {:}'.format(channels)
+    flop_A = self.conv_a.get_flops([channels[0], channels[1]])
+    flop_B = self.conv_b.get_flops([channels[1], channels[2]])
+    if hasattr(self.downsample, 'get_flops'):
+      flop_C = self.downsample.get_flops([channels[0], channels[-1]])
+    else:
+      flop_C = 0
+    if channels[0] != channels[-1] and self.downsample is None: # this short-cut will be added during the infer-train
+      flop_C = channels[0] * channels[-1] * self.conv_b.OutShape[0] * self.conv_b.OutShape[1]
+    return flop_A + flop_B + flop_C
+
+  def forward(self, inputs):
+    if self.search_mode == 'basic'   : return self.basic_forward(inputs)
+    elif self.search_mode == 'search': return self.search_forward(inputs)
+    else: raise ValueError('invalid search_mode = {:}'.format(self.search_mode))
+
+  def search_forward(self, tuple_inputs):
+    assert isinstance(tuple_inputs, tuple) and len(tuple_inputs) == 5, 'invalid type input : {:}'.format( type(tuple_inputs) )
+    inputs, expected_inC, probability, indexes, probs = tuple_inputs
+    assert indexes.size(0) == 2 and probs.size(0) == 2 and probability.size(0) == 2
+    out_a, expected_inC_a, expected_flop_a = self.conv_a( (inputs, expected_inC  , probability[0], indexes[0], probs[0]) )
+    out_b, expected_inC_b, expected_flop_b = self.conv_b( (out_a , expected_inC_a, probability[1], indexes[1], probs[1]) )
+    if self.downsample is not None:
+      residual, _, expected_flop_c = self.downsample( (inputs, expected_inC  , probability[1], indexes[1], probs[1]) )
+    else:
+      residual, expected_flop_c = inputs, 0
+    out = additive_func(residual, out_b)
+    return out, expected_inC_b, sum([expected_flop_a, expected_flop_b, expected_flop_c])
+
+  def basic_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
+    out = additive_func(residual, basicblock)
+    return nn.functional.relu(out, inplace=True)
+
+
+
+class ResNetBottleneck(nn.Module):
+  expansion = 4
+  num_conv  = 3
+  def __init__(self, inplanes, planes, stride):
+    super(ResNetBottleneck, self).__init__()
+    assert stride == 1 or stride == 2, 'invalid stride {:}'.format(stride)
+    self.conv_1x1 = ConvBNReLU(inplanes, planes, 1,      1, 0, False, has_avg=False, has_bn=True, has_relu=True)
+    self.conv_3x3 = ConvBNReLU(  planes, planes, 3, stride, 1, False, has_avg=False, has_bn=True, has_relu=True)
+    self.conv_1x4 = ConvBNReLU(planes, planes*self.expansion, 1, 1, 0, False, has_avg=False, has_bn=True, has_relu=False)
+    if stride == 2:
+      self.downsample = ConvBNReLU(inplanes, planes*self.expansion, 1, 1, 0, False, has_avg=True, has_bn=False, has_relu=False)
+    elif inplanes != planes*self.expansion:
+      self.downsample = ConvBNReLU(inplanes, planes*self.expansion, 1, 1, 0, False, has_avg=False,has_bn=True , has_relu=False)
+    else:
+      self.downsample = None
+    self.out_dim     = planes * self.expansion
+    self.search_mode = 'basic'
+
+  def get_range(self):
+    return self.conv_1x1.get_range() + self.conv_3x3.get_range() + self.conv_1x4.get_range()
+
+  def get_flops(self, channels):
+    assert len(channels) == 4, 'invalid channels : {:}'.format(channels)
+    flop_A = self.conv_1x1.get_flops([channels[0], channels[1]])
+    flop_B = self.conv_3x3.get_flops([channels[1], channels[2]])
+    flop_C = self.conv_1x4.get_flops([channels[2], channels[3]])
+    if hasattr(self.downsample, 'get_flops'):
+      flop_D = self.downsample.get_flops([channels[0], channels[-1]])
+    else:
+      flop_D = 0
+    if channels[0] != channels[-1] and self.downsample is None: # this short-cut will be added during the infer-train
+      flop_D = channels[0] * channels[-1] * self.conv_1x4.OutShape[0] * self.conv_1x4.OutShape[1]
+    return flop_A + flop_B + flop_C + flop_D
+
+  def forward(self, inputs):
+    if self.search_mode == 'basic'   : return self.basic_forward(inputs)
+    elif self.search_mode == 'search': return self.search_forward(inputs)
+    else: raise ValueError('invalid search_mode = {:}'.format(self.search_mode))
+
+  def basic_forward(self, inputs):
+    bottleneck = self.conv_1x1(inputs)
+    bottleneck = self.conv_3x3(bottleneck)
+    bottleneck = self.conv_1x4(bottleneck)
+    if self.downsample is not None: residual = self.downsample(inputs)
+    else                          : residual = inputs
+    out = additive_func(residual, bottleneck)
+    return nn.functional.relu(out, inplace=True)
+
+  def search_forward(self, tuple_inputs):
+    assert isinstance(tuple_inputs, tuple) and len(tuple_inputs) == 5, 'invalid type input : {:}'.format( type(tuple_inputs) )
+    inputs, expected_inC, probability, indexes, probs = tuple_inputs
+    assert indexes.size(0) == 3 and probs.size(0) == 3 and probability.size(0) == 3
+    out_1x1, expected_inC_1x1, expected_flop_1x1 = self.conv_1x1( (inputs, expected_inC    , probability[0], indexes[0], probs[0]) )
+    out_3x3, expected_inC_3x3, expected_flop_3x3 = self.conv_3x3( (out_1x1,expected_inC_1x1, probability[1], indexes[1], probs[1]) )
+    out_1x4, expected_inC_1x4, expected_flop_1x4 = self.conv_1x4( (out_3x3,expected_inC_3x3, probability[2], indexes[2], probs[2]) )
+    if self.downsample is not None:
+      residual, _, expected_flop_c = self.downsample( (inputs, expected_inC  , probability[2], indexes[2], probs[2]) )
+    else:
+      residual, expected_flop_c = inputs, 0
+    out = additive_func(residual, out_1x4)
+    return out, expected_inC_1x4, sum([expected_flop_1x1, expected_flop_3x3, expected_flop_1x4, expected_flop_c])
+
+
+
+class SearchShapeCifarResNet(nn.Module):
+
+  def __init__(self, block_name, depth, num_classes):
+    super(SearchShapeCifarResNet, self).__init__()
+
+    #Model type specifies number of layers for CIFAR-10 and CIFAR-100 model
+    if block_name == 'ResNetBasicblock':
+      block = ResNetBasicblock
+      assert (depth - 2) % 6 == 0, 'depth should be one of 20, 32, 44, 56, 110'
+      layer_blocks = (depth - 2) // 6
+    elif block_name == 'ResNetBottleneck':
+      block = ResNetBottleneck
+      assert (depth - 2) % 9 == 0, 'depth should be one of 164'
+      layer_blocks = (depth - 2) // 9
+    else:
+      raise ValueError('invalid block : {:}'.format(block_name))
+
+    self.message      = 'SearchShapeCifarResNet : Depth : {:} , Layers for each block : {:}'.format(depth, layer_blocks)
+    self.num_classes  = num_classes
+    self.channels     = [16]
+    self.layers       = nn.ModuleList( [ ConvBNReLU(3, 16, 3, 1, 1, False, has_avg=False, has_bn=True, has_relu=True) ] )
+    self.InShape      = None
+    self.depth_info   = OrderedDict()
+    self.depth_at_i   = OrderedDict()
+    for stage in range(3):
+      cur_block_choices = get_depth_choices(layer_blocks, False)
+      assert cur_block_choices[-1] == layer_blocks, 'stage={:}, {:} vs {:}'.format(stage, cur_block_choices, layer_blocks)
+      self.message += "\nstage={:} ::: depth-block-choices={:} for {:} blocks.".format(stage, cur_block_choices, layer_blocks)
+      block_choices, xstart = [], len(self.layers)
+      for iL in range(layer_blocks):
+        iC     = self.channels[-1]
+        planes = 16 * (2**stage)
+        stride = 2 if stage > 0 and iL == 0 else 1
+        module = block(iC, planes, stride)
+        self.channels.append( module.out_dim )
+        self.layers.append  ( module )
+        self.message += "\nstage={:}, ilayer={:02d}/{:02d}, block={:03d}, iC={:3d}, oC={:3d}, stride={:}".format(stage, iL, layer_blocks, len(self.layers)-1, iC, module.out_dim, stride)
+        # added for depth
+        layer_index = len(self.layers) - 1
+        if iL + 1 in cur_block_choices: block_choices.append( layer_index )
+        if iL + 1 == layer_blocks:
+          self.depth_info[layer_index] = {'choices': block_choices,
+                                          'stage'  : stage,
+                                          'xstart' : xstart}
+    self.depth_info_list = []
+    for xend, info in self.depth_info.items():
+      self.depth_info_list.append( (xend, info) )
+      xstart, xstage = info['xstart'], info['stage']
+      for ilayer in range(xstart, xend+1):
+        idx = bisect_right(info['choices'], ilayer-1)
+        self.depth_at_i[ilayer] = (xstage, idx)
+
+    self.avgpool     = nn.AvgPool2d(8)
+    self.classifier  = nn.Linear(module.out_dim, num_classes)
+    self.InShape     = None
+    self.tau         = -1
+    self.search_mode = 'basic'
+    #assert sum(x.num_conv for x in self.layers) + 1 == depth, 'invalid depth check {:} vs {:}'.format(sum(x.num_conv for x in self.layers)+1, depth)
+    
+    # parameters for width
+    self.Ranges = []
+    self.layer2indexRange = []
+    for i, layer in enumerate(self.layers):
+      start_index = len(self.Ranges)
+      self.Ranges += layer.get_range()
+      self.layer2indexRange.append( (start_index, len(self.Ranges)) )
+    assert len(self.Ranges) + 1 == depth, 'invalid depth check {:} vs {:}'.format(len(self.Ranges) + 1, depth)
+
+    self.register_parameter('width_attentions', nn.Parameter(torch.Tensor(len(self.Ranges), get_width_choices(None))))
+    self.register_parameter('depth_attentions', nn.Parameter(torch.Tensor(3, get_depth_choices(layer_blocks, True))))
+    nn.init.normal_(self.width_attentions, 0, 0.01)
+    nn.init.normal_(self.depth_attentions, 0, 0.01)
+    self.apply(initialize_resnet)
+
+  def arch_parameters(self, LR=None):
+    if LR is None:
+      return [self.width_attentions, self.depth_attentions]
+    else:
+      return [
+               {"params": self.width_attentions, "lr": LR},
+               {"params": self.depth_attentions, "lr": LR},
+             ]
+
+  def base_parameters(self):
+    return list(self.layers.parameters()) + list(self.avgpool.parameters()) + list(self.classifier.parameters())
+
+  def get_flop(self, mode, config_dict, extra_info):
+    if config_dict is not None: config_dict = config_dict.copy()
+    # select channels 
+    channels = [3]
+    for i, weight in enumerate(self.width_attentions):
+      if mode == 'genotype':
+        with torch.no_grad():
+          probe = nn.functional.softmax(weight, dim=0)
+          C = self.Ranges[i][ torch.argmax(probe).item() ]
+      elif mode == 'max':
+        C = self.Ranges[i][-1]
+      elif mode == 'fix':
+        C = int( math.sqrt( extra_info ) * self.Ranges[i][-1] )
+      elif mode == 'random':
+        assert isinstance(extra_info, float), 'invalid extra_info : {:}'.format(extra_info)
+        with torch.no_grad():
+          prob = nn.functional.softmax(weight, dim=0)
+          approximate_C = int( math.sqrt( extra_info ) * self.Ranges[i][-1] )
+          for j in range(prob.size(0)):
+            prob[j] = 1 / (abs(j - (approximate_C-self.Ranges[i][j])) + 0.2)
+          C = self.Ranges[i][ torch.multinomial(prob, 1, False).item() ]
+      else:
+        raise ValueError('invalid mode : {:}'.format(mode))
+      channels.append( C )
+    # select depth
+    if mode == 'genotype':
+      with torch.no_grad():
+        depth_probs = nn.functional.softmax(self.depth_attentions, dim=1)
+        choices = torch.argmax(depth_probs, dim=1).cpu().tolist()
+    elif mode == 'max' or mode == 'fix':
+      choices = [depth_probs.size(1)-1 for _ in range(depth_probs.size(0))]
+    elif mode == 'random':
+      with torch.no_grad():
+        depth_probs = nn.functional.softmax(self.depth_attentions, dim=1)
+        choices = torch.multinomial(depth_probs, 1, False).cpu().tolist()
+    else:
+      raise ValueError('invalid mode : {:}'.format(mode))
+    selected_layers = []
+    for choice, xvalue in zip(choices, self.depth_info_list):
+      xtemp = xvalue[1]['choices'][choice] - xvalue[1]['xstart'] + 1
+      selected_layers.append(xtemp)
+    flop = 0
+    for i, layer in enumerate(self.layers):
+      s, e = self.layer2indexRange[i]
+      xchl = tuple( channels[s:e+1] )
+      if i in self.depth_at_i:
+        xstagei, xatti = self.depth_at_i[i]
+        if xatti <= choices[xstagei]: # leave this depth
+          flop+= layer.get_flops(xchl)
+        else:
+          flop+= 0 # do not use this layer
+      else:
+        flop+= layer.get_flops(xchl)
+    # the last fc layer
+    flop += channels[-1] * self.classifier.out_features
+    if config_dict is None:
+      return flop / 1e6
+    else:
+      config_dict['xchannels']  = channels
+      config_dict['xblocks']    = selected_layers
+      config_dict['super_type'] = 'infer-shape'
+      config_dict['estimated_FLOP'] = flop / 1e6
+      return flop / 1e6, config_dict
+
+  def get_arch_info(self):
+    string = "for depth and width, there are {:} + {:} attention probabilities.".format(len(self.depth_attentions), len(self.width_attentions))
+    string+= '\n{:}'.format(self.depth_info)
+    discrepancy = []
+    with torch.no_grad():
+      for i, att in enumerate(self.depth_attentions):
+        prob = nn.functional.softmax(att, dim=0)
+        prob = prob.cpu() ; selc = prob.argmax().item() ; prob = prob.tolist()
+        prob = ['{:.3f}'.format(x) for x in prob]
+        xstring = '{:03d}/{:03d}-th : {:}'.format(i, len(self.depth_attentions), ' '.join(prob))
+        logt = ['{:.4f}'.format(x) for x in att.cpu().tolist()]
+        xstring += '  ||  {:17s}'.format(' '.join(logt))
+        prob = sorted( [float(x) for x in prob] )
+        disc = prob[-1] - prob[-2]
+        xstring += '  || discrepancy={:.2f} || select={:}/{:}'.format(disc, selc, len(prob))
+        discrepancy.append( disc )
+        string += '\n{:}'.format(xstring)
+      string += '\n-----------------------------------------------'
+      for i, att in enumerate(self.width_attentions):
+        prob = nn.functional.softmax(att, dim=0)
+        prob = prob.cpu() ; selc = prob.argmax().item() ; prob = prob.tolist()
+        prob = ['{:.3f}'.format(x) for x in prob]
+        xstring = '{:03d}/{:03d}-th : {:}'.format(i, len(self.width_attentions), ' '.join(prob))
+        logt = ['{:.3f}'.format(x) for x in att.cpu().tolist()]
+        xstring += '  ||  {:52s}'.format(' '.join(logt))
+        prob = sorted( [float(x) for x in prob] )
+        disc = prob[-1] - prob[-2]
+        xstring += '  || dis={:.2f} || select={:}/{:}'.format(disc, selc, len(prob))
+        discrepancy.append( disc )
+        string += '\n{:}'.format(xstring)
+    return string, discrepancy
+
+  def set_tau(self, tau_max, tau_min, epoch_ratio):
+    assert epoch_ratio >= 0 and epoch_ratio <= 1, 'invalid epoch-ratio : {:}'.format(epoch_ratio)
+    tau = tau_min + (tau_max-tau_min) * (1 + math.cos(math.pi * epoch_ratio)) / 2
+    self.tau = tau
+
+  def get_message(self):
+    return self.message
+
+  def forward(self, inputs):
+    if self.search_mode == 'basic':
+      return self.basic_forward(inputs)
+    elif self.search_mode == 'search':
+      return self.search_forward(inputs)
+    else:
+      raise ValueError('invalid search_mode = {:}'.format(self.search_mode))
+
+  def search_forward(self, inputs):
+    flop_width_probs = nn.functional.softmax(self.width_attentions, dim=1)
+    flop_depth_probs = nn.functional.softmax(self.depth_attentions, dim=1)
+    flop_depth_probs = torch.flip( torch.cumsum( torch.flip(flop_depth_probs, [1]), 1 ), [1] )
+    selected_widths, selected_width_probs = select2withP(self.width_attentions, self.tau)
+    selected_depth_probs = select2withP(self.depth_attentions, self.tau, True)
+    with torch.no_grad():
+      selected_widths = selected_widths.cpu()
+
+    x, last_channel_idx, expected_inC, flops = inputs, 0, 3, []
+    feature_maps = []
+    for i, layer in enumerate(self.layers):
+      selected_w_index = selected_widths     [last_channel_idx: last_channel_idx+layer.num_conv]
+      selected_w_probs = selected_width_probs[last_channel_idx: last_channel_idx+layer.num_conv]
+      layer_prob       = flop_width_probs    [last_channel_idx: last_channel_idx+layer.num_conv]
+      x, expected_inC, expected_flop = layer( (x, expected_inC, layer_prob, selected_w_index, selected_w_probs) )
+      feature_maps.append( x )
+      last_channel_idx += layer.num_conv
+      if i in self.depth_info: # aggregate the information
+        choices = self.depth_info[i]['choices']
+        xstagei = self.depth_info[i]['stage']
+        #print ('iL={:}, choices={:}, stage={:}, probs={:}'.format(i, choices, xstagei, selected_depth_probs[xstagei].cpu().tolist()))
+        #for A, W in zip(choices, selected_depth_probs[xstagei]):
+        #  print('Size = {:}, W = {:}'.format(feature_maps[A].size(), W))
+        possible_tensors = []
+        max_C = max( feature_maps[A].size(1) for A in choices )
+        for tempi, A in enumerate(choices):
+          xtensor = ChannelWiseInter(feature_maps[A], max_C)
+          #drop_ratio = 1-(tempi+1.0)/len(choices)
+          #xtensor = drop_path(xtensor, drop_ratio)
+          possible_tensors.append( xtensor )
+        weighted_sum = sum( xtensor * W for xtensor, W in zip(possible_tensors, selected_depth_probs[xstagei]) )
+        x = weighted_sum
+        
+      if i in self.depth_at_i:
+        xstagei, xatti = self.depth_at_i[i]
+        x_expected_flop = flop_depth_probs[xstagei, xatti] * expected_flop
+      else:
+        x_expected_flop = expected_flop
+      flops.append( x_expected_flop )
+    flops.append( expected_inC * (self.classifier.out_features*1.0/1e6) )
+    features = self.avgpool(x)
+    features = features.view(features.size(0), -1)
+    logits   = linear_forward(features, self.classifier)
+    return logits, torch.stack( [sum(flops)] )
+
+  def basic_forward(self, inputs):
+    if self.InShape is None: self.InShape = (inputs.size(-2), inputs.size(-1))
+    x = inputs
+    for i, layer in enumerate(self.layers):
+      x = layer( x )
+    features = self.avgpool(x)
+    features = features.view(features.size(0), -1)
+    logits   = self.classifier(features)
+    return features, logits

lib/models/searchs/SearchCifarResNet_depth.py

@@ -0,0 +1,337 @@
+import math, torch
+from collections import OrderedDict
+from bisect import bisect_right
+import torch.nn as nn
+from ..initialization import initialize_resnet
+from ..SharedUtils    import additive_func
+from .SoftSelect      import select2withP, ChannelWiseInter
+from .SoftSelect      import linear_forward
+from .SoftSelect      import get_width_choices
+
+
+def get_depth_choices(nDepth, return_num):
+  if nDepth == 2:
+    choices = (1, 2)
+  elif nDepth == 3:
+    choices = (1, 2, 3)
+  elif nDepth > 3:
+    choices = list(range(1, nDepth+1, 2))
+    if choices[-1] < nDepth: choices.append(nDepth)
+  else:
+    raise ValueError('invalid nDepth : {:}'.format(nDepth))
+  if return_num: return len(choices)
+  else         : return choices
+
+
+
+class ConvBNReLU(nn.Module):
+  num_conv  = 1
+  def __init__(self, nIn, nOut, kernel, stride, padding, bias, has_avg, has_bn, has_relu):
+    super(ConvBNReLU, self).__init__()
+    self.InShape  = None
+    self.OutShape = None
+    self.choices  = get_width_choices(nOut)
+    self.register_buffer('choices_tensor', torch.Tensor( self.choices ))
+
+    if has_avg : self.avg = nn.AvgPool2d(kernel_size=2, stride=2, padding=0)
+    else       : self.avg = None
+    self.conv = nn.Conv2d(nIn, nOut, kernel_size=kernel, stride=stride, padding=padding, dilation=1, groups=1, bias=bias)
+    if has_bn  : self.bn  = nn.BatchNorm2d(nOut)
+    else       : self.bn  = None
+    if has_relu: self.relu = nn.ReLU(inplace=False)
+    else       : self.relu = None
+    self.in_dim   = nIn
+    self.out_dim  = nOut
+
+  def get_flops(self, divide=1):
+    iC, oC = self.in_dim, self.out_dim
+    assert iC <= self.conv.in_channels and oC <= self.conv.out_channels, '{:} vs {:}  |  {:} vs {:}'.format(iC, self.conv.in_channels, oC, self.conv.out_channels)
+    assert isinstance(self.InShape, tuple) and len(self.InShape) == 2, 'invalid in-shape : {:}'.format(self.InShape)
+    assert isinstance(self.OutShape, tuple) and len(self.OutShape) == 2, 'invalid out-shape : {:}'.format(self.OutShape)
+    #conv_per_position_flops = self.conv.kernel_size[0] * self.conv.kernel_size[1] * iC * oC / self.conv.groups
+    conv_per_position_flops = (self.conv.kernel_size[0] * self.conv.kernel_size[1] * 1.0 / self.conv.groups)
+    all_positions = self.OutShape[0] * self.OutShape[1]
+    flops = (conv_per_position_flops * all_positions / divide) * iC * oC
+    if self.conv.bias is not None: flops += all_positions / divide
+    return flops
+
+  def forward(self, inputs):
+    if self.avg : out = self.avg( inputs )
+    else        : out = inputs
+    conv = self.conv( out )
+    if self.bn  : out = self.bn( conv )
+    else        : out = conv
+    if self.relu: out = self.relu( out )
+    else        : out = out
+    if self.InShape is None:
+      self.InShape  = (inputs.size(-2), inputs.size(-1))
+      self.OutShape = (out.size(-2)   , out.size(-1))
+    return out
+
+
+class ResNetBasicblock(nn.Module):
+  expansion = 1
+  num_conv  = 2
+  def __init__(self, inplanes, planes, stride):
+    super(ResNetBasicblock, self).__init__()
+    assert stride == 1 or stride == 2, 'invalid stride {:}'.format(stride)
+    self.conv_a = ConvBNReLU(inplanes, planes, 3, stride, 1, False, has_avg=False, has_bn=True, has_relu=True)
+    self.conv_b = ConvBNReLU(  planes, planes, 3,      1, 1, False, has_avg=False, has_bn=True, has_relu=False)
+    if stride == 2:
+      self.downsample = ConvBNReLU(inplanes, planes, 1, 1, 0, False, has_avg=True, has_bn=False, has_relu=False)
+    elif inplanes != planes:
+      self.downsample = ConvBNReLU(inplanes, planes, 1, 1, 0, False, has_avg=False,has_bn=True , has_relu=False)
+    else:
+      self.downsample = None
+    self.out_dim     = planes
+    self.search_mode = 'basic'
+
+  def get_flops(self, divide=1):
+    flop_A = self.conv_a.get_flops(divide)
+    flop_B = self.conv_b.get_flops(divide)
+    if hasattr(self.downsample, 'get_flops'):
+      flop_C = self.downsample.get_flops(divide)
+    else:
+      flop_C = 0
+    return flop_A + flop_B + flop_C
+
+  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
+    out = additive_func(residual, basicblock)
+    return nn.functional.relu(out, inplace=True)
+
+
+
+class ResNetBottleneck(nn.Module):
+  expansion = 4
+  num_conv  = 3
+  def __init__(self, inplanes, planes, stride):
+    super(ResNetBottleneck, self).__init__()
+    assert stride == 1 or stride == 2, 'invalid stride {:}'.format(stride)
+    self.conv_1x1 = ConvBNReLU(inplanes, planes, 1,      1, 0, False, has_avg=False, has_bn=True, has_relu=True)
+    self.conv_3x3 = ConvBNReLU(  planes, planes, 3, stride, 1, False, has_avg=False, has_bn=True, has_relu=True)
+    self.conv_1x4 = ConvBNReLU(planes, planes*self.expansion, 1, 1, 0, False, has_avg=False, has_bn=True, has_relu=False)
+    if stride == 2:
+      self.downsample = ConvBNReLU(inplanes, planes*self.expansion, 1, 1, 0, False, has_avg=True, has_bn=False, has_relu=False)
+    elif inplanes != planes*self.expansion:
+      self.downsample = ConvBNReLU(inplanes, planes*self.expansion, 1, 1, 0, False, has_avg=False,has_bn=True , has_relu=False)
+    else:
+      self.downsample = None
+    self.out_dim     = planes * self.expansion
+    self.search_mode = 'basic'
+
+  def get_range(self):
+    return self.conv_1x1.get_range() + self.conv_3x3.get_range() + self.conv_1x4.get_range()
+
+  def get_flops(self, divide):
+    flop_A = self.conv_1x1.get_flops(divide)
+    flop_B = self.conv_3x3.get_flops(divide)
+    flop_C = self.conv_1x4.get_flops(divide)
+    if hasattr(self.downsample, 'get_flops'):
+      flop_D = self.downsample.get_flops(divide)
+    else:
+      flop_D = 0
+    return flop_A + flop_B + flop_C + flop_D
+
+  def forward(self, inputs):
+    bottleneck = self.conv_1x1(inputs)
+    bottleneck = self.conv_3x3(bottleneck)
+    bottleneck = self.conv_1x4(bottleneck)
+    if self.downsample is not None: residual = self.downsample(inputs)
+    else                          : residual = inputs
+    out = additive_func(residual, bottleneck)
+    return nn.functional.relu(out, inplace=True)
+
+
+class SearchDepthCifarResNet(nn.Module):
+
+  def __init__(self, block_name, depth, num_classes):
+    super(SearchDepthCifarResNet, self).__init__()
+
+    #Model type specifies number of layers for CIFAR-10 and CIFAR-100 model
+    if block_name == 'ResNetBasicblock':
+      block = ResNetBasicblock
+      assert (depth - 2) % 6 == 0, 'depth should be one of 20, 32, 44, 56, 110'
+      layer_blocks = (depth - 2) // 6
+    elif block_name == 'ResNetBottleneck':
+      block = ResNetBottleneck
+      assert (depth - 2) % 9 == 0, 'depth should be one of 164'
+      layer_blocks = (depth - 2) // 9
+    else:
+      raise ValueError('invalid block : {:}'.format(block_name))
+
+    self.message      = 'SearchShapeCifarResNet : Depth : {:} , Layers for each block : {:}'.format(depth, layer_blocks)
+    self.num_classes  = num_classes
+    self.channels     = [16]
+    self.layers       = nn.ModuleList( [ ConvBNReLU(3, 16, 3, 1, 1, False, has_avg=False, has_bn=True, has_relu=True) ] )
+    self.InShape      = None
+    self.depth_info   = OrderedDict()
+    self.depth_at_i   = OrderedDict()
+    for stage in range(3):
+      cur_block_choices = get_depth_choices(layer_blocks, False)
+      assert cur_block_choices[-1] == layer_blocks, 'stage={:}, {:} vs {:}'.format(stage, cur_block_choices, layer_blocks)
+      self.message += "\nstage={:} ::: depth-block-choices={:} for {:} blocks.".format(stage, cur_block_choices, layer_blocks)
+      block_choices, xstart = [], len(self.layers)
+      for iL in range(layer_blocks):
+        iC     = self.channels[-1]
+        planes = 16 * (2**stage)
+        stride = 2 if stage > 0 and iL == 0 else 1
+        module = block(iC, planes, stride)
+        self.channels.append( module.out_dim )
+        self.layers.append  ( module )
+        self.message += "\nstage={:}, ilayer={:02d}/{:02d}, block={:03d}, iC={:3d}, oC={:3d}, stride={:}".format(stage, iL, layer_blocks, len(self.layers)-1, iC, module.out_dim, stride)
+        # added for depth
+        layer_index = len(self.layers) - 1
+        if iL + 1 in cur_block_choices: block_choices.append( layer_index )
+        if iL + 1 == layer_blocks:
+          self.depth_info[layer_index] = {'choices': block_choices,
+                                          'stage'  : stage,
+                                          'xstart' : xstart}
+    self.depth_info_list = []
+    for xend, info in self.depth_info.items():
+      self.depth_info_list.append( (xend, info) )
+      xstart, xstage = info['xstart'], info['stage']
+      for ilayer in range(xstart, xend+1):
+        idx = bisect_right(info['choices'], ilayer-1)
+        self.depth_at_i[ilayer] = (xstage, idx)
+
+    self.avgpool     = nn.AvgPool2d(8)
+    self.classifier  = nn.Linear(module.out_dim, num_classes)
+    self.InShape     = None
+    self.tau         = -1
+    self.search_mode = 'basic'
+    #assert sum(x.num_conv for x in self.layers) + 1 == depth, 'invalid depth check {:} vs {:}'.format(sum(x.num_conv for x in self.layers)+1, depth)
+    
+
+    self.register_parameter('depth_attentions', nn.Parameter(torch.Tensor(3, get_depth_choices(layer_blocks, True))))
+    nn.init.normal_(self.depth_attentions, 0, 0.01)
+    self.apply(initialize_resnet)
+
+  def arch_parameters(self):
+    return [self.depth_attentions]
+
+  def base_parameters(self):
+    return list(self.layers.parameters()) + list(self.avgpool.parameters()) + list(self.classifier.parameters())
+
+  def get_flop(self, mode, config_dict, extra_info):
+    if config_dict is not None: config_dict = config_dict.copy()
+    # select depth
+    if mode == 'genotype':
+      with torch.no_grad():
+        depth_probs = nn.functional.softmax(self.depth_attentions, dim=1)
+        choices = torch.argmax(depth_probs, dim=1).cpu().tolist()
+    elif mode == 'max':
+      choices = [depth_probs.size(1)-1 for _ in range(depth_probs.size(0))]
+    elif mode == 'random':
+      with torch.no_grad():
+        depth_probs = nn.functional.softmax(self.depth_attentions, dim=1)
+        choices = torch.multinomial(depth_probs, 1, False).cpu().tolist()
+    else:
+      raise ValueError('invalid mode : {:}'.format(mode))
+    selected_layers = []
+    for choice, xvalue in zip(choices, self.depth_info_list):
+      xtemp = xvalue[1]['choices'][choice] - xvalue[1]['xstart'] + 1
+      selected_layers.append(xtemp)
+    flop = 0
+    for i, layer in enumerate(self.layers):
+      if i in self.depth_at_i:
+        xstagei, xatti = self.depth_at_i[i]
+        if xatti <= choices[xstagei]: # leave this depth
+          flop+= layer.get_flops()
+        else:
+          flop+= 0 # do not use this layer
+      else:
+        flop+= layer.get_flops()
+    # the last fc layer
+    flop += self.classifier.in_features * self.classifier.out_features
+    if config_dict is None:
+      return flop / 1e6
+    else:
+      config_dict['xblocks']    = selected_layers
+      config_dict['super_type'] = 'infer-depth'
+      config_dict['estimated_FLOP'] = flop / 1e6
+      return flop / 1e6, config_dict
+
+  def get_arch_info(self):
+    string = "for depth, there are {:} attention probabilities.".format(len(self.depth_attentions))
+    string+= '\n{:}'.format(self.depth_info)
+    discrepancy = []
+    with torch.no_grad():
+      for i, att in enumerate(self.depth_attentions):
+        prob = nn.functional.softmax(att, dim=0)
+        prob = prob.cpu() ; selc = prob.argmax().item() ; prob = prob.tolist()
+        prob = ['{:.3f}'.format(x) for x in prob]
+        xstring = '{:03d}/{:03d}-th : {:}'.format(i, len(self.depth_attentions), ' '.join(prob))
+        logt = ['{:.4f}'.format(x) for x in att.cpu().tolist()]
+        xstring += '  ||  {:17s}'.format(' '.join(logt))
+        prob = sorted( [float(x) for x in prob] )
+        disc = prob[-1] - prob[-2]
+        xstring += '  || discrepancy={:.2f} || select={:}/{:}'.format(disc, selc, len(prob))
+        discrepancy.append( disc )
+        string += '\n{:}'.format(xstring)
+    return string, discrepancy
+
+  def set_tau(self, tau_max, tau_min, epoch_ratio):
+    assert epoch_ratio >= 0 and epoch_ratio <= 1, 'invalid epoch-ratio : {:}'.format(epoch_ratio)
+    tau = tau_min + (tau_max-tau_min) * (1 + math.cos(math.pi * epoch_ratio)) / 2
+    self.tau = tau
+
+  def get_message(self):
+    return self.message
+
+  def forward(self, inputs):
+    if self.search_mode == 'basic':
+      return self.basic_forward(inputs)
+    elif self.search_mode == 'search':
+      return self.search_forward(inputs)
+    else:
+      raise ValueError('invalid search_mode = {:}'.format(self.search_mode))
+
+  def search_forward(self, inputs):
+    flop_depth_probs = nn.functional.softmax(self.depth_attentions, dim=1)
+    flop_depth_probs = torch.flip( torch.cumsum( torch.flip(flop_depth_probs, [1]), 1 ), [1] )
+    selected_depth_probs = select2withP(self.depth_attentions, self.tau, True)
+
+    x, flops = inputs, []
+    feature_maps = []
+    for i, layer in enumerate(self.layers):
+      layer_i = layer( x )
+      feature_maps.append( layer_i )
+      if i in self.depth_info: # aggregate the information
+        choices = self.depth_info[i]['choices']
+        xstagei = self.depth_info[i]['stage']
+        possible_tensors = []
+        for tempi, A in enumerate(choices):
+          xtensor = feature_maps[A]
+          possible_tensors.append( xtensor )
+        weighted_sum = sum( xtensor * W for xtensor, W in zip(possible_tensors, selected_depth_probs[xstagei]) )
+        x = weighted_sum
+      else:
+        x = layer_i
+       
+      if i in self.depth_at_i:
+        xstagei, xatti = self.depth_at_i[i]
+        #print ('layer-{:03d}, stage={:}, att={:}, prob={:}, flop={:}'.format(i, xstagei, xatti, flop_depth_probs[xstagei, xatti].item(), layer.get_flops(1e6)))
+        x_expected_flop = flop_depth_probs[xstagei, xatti] * layer.get_flops(1e6)
+      else:
+        x_expected_flop = layer.get_flops(1e6)
+      flops.append( x_expected_flop )
+    flops.append( (self.classifier.in_features * self.classifier.out_features*1.0/1e6) )
+
+    features = self.avgpool(x)
+    features = features.view(features.size(0), -1)
+    logits   = linear_forward(features, self.classifier)
+    return logits, torch.stack( [sum(flops)] )
+
+  def basic_forward(self, inputs):
+    if self.InShape is None: self.InShape = (inputs.size(-2), inputs.size(-1))
+    x = inputs
+    for i, layer in enumerate(self.layers):
+      x = layer( x )
+    features = self.avgpool(x)
+    features = features.view(features.size(0), -1)
+    logits   = self.classifier(features)
+    return features, logits

lib/models/searchs/SearchCifarResNet_width.py

@@ -0,0 +1,391 @@
+import math, torch
+import torch.nn as nn
+from ..initialization import initialize_resnet
+from ..SharedUtils    import additive_func
+from .SoftSelect      import select2withP, ChannelWiseInter
+from .SoftSelect      import linear_forward
+from .SoftSelect      import get_width_choices as get_choices
+
+
+def conv_forward(inputs, conv, choices):
+  iC = conv.in_channels
+  fill_size = list(inputs.size())
+  fill_size[1] = iC - fill_size[1]
+  filled  = torch.zeros(fill_size, device=inputs.device)
+  xinputs = torch.cat((inputs, filled), dim=1)
+  outputs = conv(xinputs)
+  selecteds = [outputs[:,:oC] for oC in choices]
+  return selecteds
+
+
+class ConvBNReLU(nn.Module):
+  num_conv  = 1
+  def __init__(self, nIn, nOut, kernel, stride, padding, bias, has_avg, has_bn, has_relu):
+    super(ConvBNReLU, self).__init__()
+    self.InShape  = None
+    self.OutShape = None
+    self.choices  = get_choices(nOut)
+    self.register_buffer('choices_tensor', torch.Tensor( self.choices ))
+
+    if has_avg : self.avg = nn.AvgPool2d(kernel_size=2, stride=2, padding=0)
+    else       : self.avg = None
+    self.conv = nn.Conv2d(nIn, nOut, kernel_size=kernel, stride=stride, padding=padding, dilation=1, groups=1, bias=bias)
+    #if has_bn  : self.bn  = nn.BatchNorm2d(nOut)
+    #else       : self.bn  = None
+    self.has_bn = has_bn
+    self.BNs  = nn.ModuleList()
+    for i, _out in enumerate(self.choices):
+      self.BNs.append(nn.BatchNorm2d(_out))
+    if has_relu: self.relu = nn.ReLU(inplace=True)
+    else       : self.relu = None
+    self.in_dim   = nIn
+    self.out_dim  = nOut
+    self.search_mode = 'basic'
+
+  def get_flops(self, channels, check_range=True, divide=1):
+    iC, oC = channels
+    if check_range: assert iC <= self.conv.in_channels and oC <= self.conv.out_channels, '{:} vs {:}  |  {:} vs {:}'.format(iC, self.conv.in_channels, oC, self.conv.out_channels)
+    assert isinstance(self.InShape, tuple) and len(self.InShape) == 2, 'invalid in-shape : {:}'.format(self.InShape)
+    assert isinstance(self.OutShape, tuple) and len(self.OutShape) == 2, 'invalid out-shape : {:}'.format(self.OutShape)
+    #conv_per_position_flops = self.conv.kernel_size[0] * self.conv.kernel_size[1] * iC * oC / self.conv.groups
+    conv_per_position_flops = (self.conv.kernel_size[0] * self.conv.kernel_size[1] * 1.0 / self.conv.groups)
+    all_positions = self.OutShape[0] * self.OutShape[1]
+    flops = (conv_per_position_flops * all_positions / divide) * iC * oC
+    if self.conv.bias is not None: flops += all_positions / divide
+    return flops
+
+  def get_range(self):
+    return [self.choices]
+
+  def forward(self, inputs):
+    if self.search_mode == 'basic':
+      return self.basic_forward(inputs)
+    elif self.search_mode == 'search':
+      return self.search_forward(inputs)
+    else:
+      raise ValueError('invalid search_mode = {:}'.format(self.search_mode))
+
+  def search_forward(self, tuple_inputs):
+    assert isinstance(tuple_inputs, tuple) and len(tuple_inputs) == 5, 'invalid type input : {:}'.format( type(tuple_inputs) )
+    inputs, expected_inC, probability, index, prob = tuple_inputs
+    index, prob = torch.squeeze(index).tolist(), torch.squeeze(prob)
+    probability = torch.squeeze(probability)
+    assert len(index) == 2, 'invalid length : {:}'.format(index)
+    # compute expected flop
+    #coordinates   = torch.arange(self.x_range[0], self.x_range[1]+1).type_as(probability)
+    expected_outC = (self.choices_tensor * probability).sum()
+    expected_flop = self.get_flops([expected_inC, expected_outC], False, 1e6)
+    if self.avg : out = self.avg( inputs )
+    else        : out = inputs
+    # convolutional layer
+    out_convs = conv_forward(out, self.conv, [self.choices[i] for i in index])
+    out_bns   = [self.BNs[idx](out_conv) for idx, out_conv in zip(index, out_convs)]
+    # merge
+    out_channel = max([x.size(1) for x in out_bns])
+    outA = ChannelWiseInter(out_bns[0], out_channel)
+    outB = ChannelWiseInter(out_bns[1], out_channel)
+    out  = outA * prob[0] + outB * prob[1]
+    #out = additive_func(out_bns[0]*prob[0], out_bns[1]*prob[1])
+
+    if self.relu: out = self.relu( out )
+    else        : out = out
+    return out, expected_outC, expected_flop
+
+  def basic_forward(self, inputs):
+    if self.avg : out = self.avg( inputs )
+    else        : out = inputs
+    conv = self.conv( out )
+    if self.has_bn:out= self.BNs[-1]( conv )
+    else        : out = conv
+    if self.relu: out = self.relu( out )
+    else        : out = out
+    if self.InShape is None:
+      self.InShape  = (inputs.size(-2), inputs.size(-1))
+      self.OutShape = (out.size(-2)   , out.size(-1))
+    return out
+
+
+class ResNetBasicblock(nn.Module):
+  expansion = 1
+  num_conv  = 2
+  def __init__(self, inplanes, planes, stride):
+    super(ResNetBasicblock, self).__init__()
+    assert stride == 1 or stride == 2, 'invalid stride {:}'.format(stride)
+    self.conv_a = ConvBNReLU(inplanes, planes, 3, stride, 1, False, has_avg=False, has_bn=True, has_relu=True)
+    self.conv_b = ConvBNReLU(  planes, planes, 3,      1, 1, False, has_avg=False, has_bn=True, has_relu=False)
+    if stride == 2:
+      self.downsample = ConvBNReLU(inplanes, planes, 1, 1, 0, False, has_avg=True, has_bn=False, has_relu=False)
+    elif inplanes != planes:
+      self.downsample = ConvBNReLU(inplanes, planes, 1, 1, 0, False, has_avg=False,has_bn=True , has_relu=False)
+    else:
+      self.downsample = None
+    self.out_dim     = planes
+    self.search_mode = 'basic'
+
+  def get_range(self):
+    return self.conv_a.get_range() + self.conv_b.get_range()
+
+  def get_flops(self, channels):
+    assert len(channels) == 3, 'invalid channels : {:}'.format(channels)
+    flop_A = self.conv_a.get_flops([channels[0], channels[1]])
+    flop_B = self.conv_b.get_flops([channels[1], channels[2]])
+    if hasattr(self.downsample, 'get_flops'):
+      flop_C = self.downsample.get_flops([channels[0], channels[-1]])
+    else:
+      flop_C = 0
+    if channels[0] != channels[-1] and self.downsample is None: # this short-cut will be added during the infer-train
+      flop_C = channels[0] * channels[-1] * self.conv_b.OutShape[0] * self.conv_b.OutShape[1]
+    return flop_A + flop_B + flop_C
+
+  def forward(self, inputs):
+    if self.search_mode == 'basic'   : return self.basic_forward(inputs)
+    elif self.search_mode == 'search': return self.search_forward(inputs)
+    else: raise ValueError('invalid search_mode = {:}'.format(self.search_mode))
+
+  def search_forward(self, tuple_inputs):
+    assert isinstance(tuple_inputs, tuple) and len(tuple_inputs) == 5, 'invalid type input : {:}'.format( type(tuple_inputs) )
+    inputs, expected_inC, probability, indexes, probs = tuple_inputs
+    assert indexes.size(0) == 2 and probs.size(0) == 2 and probability.size(0) == 2
+    out_a, expected_inC_a, expected_flop_a = self.conv_a( (inputs, expected_inC  , probability[0], indexes[0], probs[0]) )
+    out_b, expected_inC_b, expected_flop_b = self.conv_b( (out_a , expected_inC_a, probability[1], indexes[1], probs[1]) )
+    if self.downsample is not None:
+      residual, _, expected_flop_c = self.downsample( (inputs, expected_inC  , probability[1], indexes[1], probs[1]) )
+    else:
+      residual, expected_flop_c = inputs, 0
+    out = additive_func(residual, out_b)
+    return out, expected_inC_b, sum([expected_flop_a, expected_flop_b, expected_flop_c])
+
+  def basic_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
+    out = additive_func(residual, basicblock)
+    return nn.functional.relu(out, inplace=True)
+
+
+
+class ResNetBottleneck(nn.Module):
+  expansion = 4
+  num_conv  = 3
+  def __init__(self, inplanes, planes, stride):
+    super(ResNetBottleneck, self).__init__()
+    assert stride == 1 or stride == 2, 'invalid stride {:}'.format(stride)
+    self.conv_1x1 = ConvBNReLU(inplanes, planes, 1,      1, 0, False, has_avg=False, has_bn=True, has_relu=True)
+    self.conv_3x3 = ConvBNReLU(  planes, planes, 3, stride, 1, False, has_avg=False, has_bn=True, has_relu=True)
+    self.conv_1x4 = ConvBNReLU(planes, planes*self.expansion, 1, 1, 0, False, has_avg=False, has_bn=True, has_relu=False)
+    if stride == 2:
+      self.downsample = ConvBNReLU(inplanes, planes*self.expansion, 1, 1, 0, False, has_avg=True, has_bn=False, has_relu=False)
+    elif inplanes != planes*self.expansion:
+      self.downsample = ConvBNReLU(inplanes, planes*self.expansion, 1, 1, 0, False, has_avg=False,has_bn=True , has_relu=False)
+    else:
+      self.downsample = None
+    self.out_dim     = planes * self.expansion
+    self.search_mode = 'basic'
+
+  def get_range(self):
+    return self.conv_1x1.get_range() + self.conv_3x3.get_range() + self.conv_1x4.get_range()
+
+  def get_flops(self, channels):
+    assert len(channels) == 4, 'invalid channels : {:}'.format(channels)
+    flop_A = self.conv_1x1.get_flops([channels[0], channels[1]])
+    flop_B = self.conv_3x3.get_flops([channels[1], channels[2]])
+    flop_C = self.conv_1x4.get_flops([channels[2], channels[3]])
+    if hasattr(self.downsample, 'get_flops'):
+      flop_D = self.downsample.get_flops([channels[0], channels[-1]])
+    else:
+      flop_D = 0
+    if channels[0] != channels[-1] and self.downsample is None: # this short-cut will be added during the infer-train
+      flop_D = channels[0] * channels[-1] * self.conv_1x4.OutShape[0] * self.conv_1x4.OutShape[1]
+    return flop_A + flop_B + flop_C + flop_D
+
+  def forward(self, inputs):
+    if self.search_mode == 'basic'   : return self.basic_forward(inputs)
+    elif self.search_mode == 'search': return self.search_forward(inputs)
+    else: raise ValueError('invalid search_mode = {:}'.format(self.search_mode))
+
+  def basic_forward(self, inputs):
+    bottleneck = self.conv_1x1(inputs)
+    bottleneck = self.conv_3x3(bottleneck)
+    bottleneck = self.conv_1x4(bottleneck)
+    if self.downsample is not None: residual = self.downsample(inputs)
+    else                          : residual = inputs
+    out = additive_func(residual, bottleneck)
+    return nn.functional.relu(out, inplace=True)
+
+  def search_forward(self, tuple_inputs):
+    assert isinstance(tuple_inputs, tuple) and len(tuple_inputs) == 5, 'invalid type input : {:}'.format( type(tuple_inputs) )
+    inputs, expected_inC, probability, indexes, probs = tuple_inputs
+    assert indexes.size(0) == 3 and probs.size(0) == 3 and probability.size(0) == 3
+    out_1x1, expected_inC_1x1, expected_flop_1x1 = self.conv_1x1( (inputs, expected_inC    , probability[0], indexes[0], probs[0]) )
+    out_3x3, expected_inC_3x3, expected_flop_3x3 = self.conv_3x3( (out_1x1,expected_inC_1x1, probability[1], indexes[1], probs[1]) )
+    out_1x4, expected_inC_1x4, expected_flop_1x4 = self.conv_1x4( (out_3x3,expected_inC_3x3, probability[2], indexes[2], probs[2]) )
+    if self.downsample is not None:
+      residual, _, expected_flop_c = self.downsample( (inputs, expected_inC  , probability[2], indexes[2], probs[2]) )
+    else:
+      residual, expected_flop_c = inputs, 0
+    out = additive_func(residual, out_1x4)
+    return out, expected_inC_1x4, sum([expected_flop_1x1, expected_flop_3x3, expected_flop_1x4, expected_flop_c])
+
+
+
+class SearchWidthCifarResNet(nn.Module):
+
+  def __init__(self, block_name, depth, num_classes):
+    super(SearchWidthCifarResNet, self).__init__()
+
+    #Model type specifies number of layers for CIFAR-10 and CIFAR-100 model
+    if block_name == 'ResNetBasicblock':
+      block = ResNetBasicblock
+      assert (depth - 2) % 6 == 0, 'depth should be one of 20, 32, 44, 56, 110'
+      layer_blocks = (depth - 2) // 6
+    elif block_name == 'ResNetBottleneck':
+      block = ResNetBottleneck
+      assert (depth - 2) % 9 == 0, 'depth should be one of 164'
+      layer_blocks = (depth - 2) // 9
+    else:
+      raise ValueError('invalid block : {:}'.format(block_name))
+
+    self.message     = 'SearchWidthCifarResNet : Depth : {:} , Layers for each block : {:}'.format(depth, layer_blocks)
+    self.num_classes = num_classes
+    self.channels    = [16]
+    self.layers      = nn.ModuleList( [ ConvBNReLU(3, 16, 3, 1, 1, False, has_avg=False, has_bn=True, has_relu=True) ] )
+    self.InShape     = None
+    for stage in range(3):
+      for iL in range(layer_blocks):
+        iC     = self.channels[-1]
+        planes = 16 * (2**stage)
+        stride = 2 if stage > 0 and iL == 0 else 1
+        module = block(iC, planes, stride)
+        self.channels.append( module.out_dim )
+        self.layers.append  ( module )
+        self.message += "\nstage={:}, ilayer={:02d}/{:02d}, block={:03d}, iC={:3d}, oC={:3d}, stride={:}".format(stage, iL, layer_blocks, len(self.layers)-1, iC, module.out_dim, stride)
+  
+    self.avgpool     = nn.AvgPool2d(8)
+    self.classifier  = nn.Linear(module.out_dim, num_classes)
+    self.InShape     = None
+    self.tau         = -1
+    self.search_mode = 'basic'
+    #assert sum(x.num_conv for x in self.layers) + 1 == depth, 'invalid depth check {:} vs {:}'.format(sum(x.num_conv for x in self.layers)+1, depth)
+    
+    # parameters for width
+    self.Ranges = []
+    self.layer2indexRange = []
+    for i, layer in enumerate(self.layers):
+      start_index = len(self.Ranges)
+      self.Ranges += layer.get_range()
+      self.layer2indexRange.append( (start_index, len(self.Ranges)) )
+    assert len(self.Ranges) + 1 == depth, 'invalid depth check {:} vs {:}'.format(len(self.Ranges) + 1, depth)
+
+    self.register_parameter('width_attentions', nn.Parameter(torch.Tensor(len(self.Ranges), get_choices(None))))
+    nn.init.normal_(self.width_attentions, 0, 0.01)
+    self.apply(initialize_resnet)
+
+  def arch_parameters(self):
+    return [self.width_attentions]
+
+  def base_parameters(self):
+    return list(self.layers.parameters()) + list(self.avgpool.parameters()) + list(self.classifier.parameters())
+
+  def get_flop(self, mode, config_dict, extra_info):
+    if config_dict is not None: config_dict = config_dict.copy()
+    #weights = [F.softmax(x, dim=0) for x in self.width_attentions]
+    channels = [3]
+    for i, weight in enumerate(self.width_attentions):
+      if mode == 'genotype':
+        with torch.no_grad():
+          probe = nn.functional.softmax(weight, dim=0)
+          C = self.Ranges[i][ torch.argmax(probe).item() ]
+      elif mode == 'max':
+        C = self.Ranges[i][-1]
+      elif mode == 'fix':
+        C = int( math.sqrt( extra_info ) * self.Ranges[i][-1] )
+      elif mode == 'random':
+        assert isinstance(extra_info, float), 'invalid extra_info : {:}'.format(extra_info)
+        with torch.no_grad():
+          prob = nn.functional.softmax(weight, dim=0)
+          approximate_C = int( math.sqrt( extra_info ) * self.Ranges[i][-1] )
+          for j in range(prob.size(0)):
+            prob[j] = 1 / (abs(j - (approximate_C-self.Ranges[i][j])) + 0.2)
+          C = self.Ranges[i][ torch.multinomial(prob, 1, False).item() ]
+      else:
+        raise ValueError('invalid mode : {:}'.format(mode))
+      channels.append( C )
+    flop = 0
+    for i, layer in enumerate(self.layers):
+      s, e = self.layer2indexRange[i]
+      xchl = tuple( channels[s:e+1] )
+      flop+= layer.get_flops(xchl)
+    # the last fc layer
+    flop += channels[-1] * self.classifier.out_features
+    if config_dict is None:
+      return flop / 1e6
+    else:
+      config_dict['xchannels']  = channels
+      config_dict['super_type'] = 'infer-width'
+      config_dict['estimated_FLOP'] = flop / 1e6
+      return flop / 1e6, config_dict
+
+  def get_arch_info(self):
+    string = "for width, there are {:} attention probabilities.".format(len(self.width_attentions))
+    discrepancy = []
+    with torch.no_grad():
+      for i, att in enumerate(self.width_attentions):
+        prob = nn.functional.softmax(att, dim=0)
+        prob = prob.cpu() ; selc = prob.argmax().item() ; prob = prob.tolist()
+        prob = ['{:.3f}'.format(x) for x in prob]
+        xstring = '{:03d}/{:03d}-th : {:}'.format(i, len(self.width_attentions), ' '.join(prob))
+        logt = ['{:.3f}'.format(x) for x in att.cpu().tolist()]
+        xstring += '  ||  {:52s}'.format(' '.join(logt))
+        prob = sorted( [float(x) for x in prob] )
+        disc = prob[-1] - prob[-2]
+        xstring += '  || dis={:.2f} || select={:}/{:}'.format(disc, selc, len(prob))
+        discrepancy.append( disc )
+        string += '\n{:}'.format(xstring)
+    return string, discrepancy
+
+  def set_tau(self, tau_max, tau_min, epoch_ratio):
+    assert epoch_ratio >= 0 and epoch_ratio <= 1, 'invalid epoch-ratio : {:}'.format(epoch_ratio)
+    tau = tau_min + (tau_max-tau_min) * (1 + math.cos(math.pi * epoch_ratio)) / 2
+    self.tau = tau
+
+  def get_message(self):
+    return self.message
+
+  def forward(self, inputs):
+    if self.search_mode == 'basic':
+      return self.basic_forward(inputs)
+    elif self.search_mode == 'search':
+      return self.search_forward(inputs)
+    else:
+      raise ValueError('invalid search_mode = {:}'.format(self.search_mode))
+
+  def search_forward(self, inputs):
+    flop_probs = nn.functional.softmax(self.width_attentions, dim=1)
+    selected_widths, selected_probs = select2withP(self.width_attentions, self.tau)
+    with torch.no_grad():
+      selected_widths = selected_widths.cpu()
+
+    x, last_channel_idx, expected_inC, flops = inputs, 0, 3, []
+    for i, layer in enumerate(self.layers):
+      selected_w_index = selected_widths[last_channel_idx: last_channel_idx+layer.num_conv]
+      selected_w_probs = selected_probs[last_channel_idx: last_channel_idx+layer.num_conv]
+      layer_prob       = flop_probs[last_channel_idx: last_channel_idx+layer.num_conv]
+      x, expected_inC, expected_flop = layer( (x, expected_inC, layer_prob, selected_w_index, selected_w_probs) )
+      last_channel_idx += layer.num_conv
+      flops.append( expected_flop )
+    flops.append( expected_inC * (self.classifier.out_features*1.0/1e6) )
+    features = self.avgpool(x)
+    features = features.view(features.size(0), -1)
+    logits   = linear_forward(features, self.classifier)
+    return logits, torch.stack( [sum(flops)] )
+
+  def basic_forward(self, inputs):
+    if self.InShape is None: self.InShape = (inputs.size(-2), inputs.size(-1))
+    x = inputs
+    for i, layer in enumerate(self.layers):
+      x = layer( x )
+    features = self.avgpool(x)
+    features = features.view(features.size(0), -1)
+    logits   = self.classifier(features)
+    return features, logits

lib/models/searchs/SearchImagenetResNet.py

@@ -0,0 +1,483 @@
+import math, torch
+from collections import OrderedDict
+from bisect import bisect_right
+import torch.nn as nn
+from ..initialization import initialize_resnet
+from ..SharedUtils    import additive_func
+from .SoftSelect      import select2withP, ChannelWiseInter
+from .SoftSelect      import linear_forward
+from .SoftSelect      import get_width_choices
+
+
+def get_depth_choices(layers):
+  min_depth = min(layers)
+  info = {'num': min_depth}
+  for i, depth in enumerate(layers):
+    choices = []
+    for j in range(1, min_depth+1):
+      choices.append( int( float(depth)*j/min_depth ) )
+    info[i] = choices
+  return info
+
+
+def conv_forward(inputs, conv, choices):
+  iC = conv.in_channels
+  fill_size = list(inputs.size())
+  fill_size[1] = iC - fill_size[1]
+  filled  = torch.zeros(fill_size, device=inputs.device)
+  xinputs = torch.cat((inputs, filled), dim=1)
+  outputs = conv(xinputs)
+  selecteds = [outputs[:,:oC] for oC in choices]
+  return selecteds
+
+
+class ConvBNReLU(nn.Module):
+  num_conv  = 1
+  def __init__(self, nIn, nOut, kernel, stride, padding, bias, has_avg, has_bn, has_relu, last_max_pool=False):
+    super(ConvBNReLU, self).__init__()
+    self.InShape  = None
+    self.OutShape = None
+    self.choices  = get_width_choices(nOut)
+    self.register_buffer('choices_tensor', torch.Tensor( self.choices ))
+
+    if has_avg : self.avg = nn.AvgPool2d(kernel_size=2, stride=2, padding=0)
+    else       : self.avg = None
+    self.conv = nn.Conv2d(nIn, nOut, kernel_size=kernel, stride=stride, padding=padding, dilation=1, groups=1, bias=bias)
+    #if has_bn  : self.bn  = nn.BatchNorm2d(nOut)
+    #else       : self.bn  = None
+    self.has_bn = has_bn
+    self.BNs  = nn.ModuleList()
+    for i, _out in enumerate(self.choices):
+      self.BNs.append(nn.BatchNorm2d(_out))
+    if has_relu: self.relu = nn.ReLU(inplace=True)
+    else       : self.relu = None
+  
+    if last_max_pool: self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+    else            : self.maxpool = None
+    self.in_dim   = nIn
+    self.out_dim  = nOut
+    self.search_mode = 'basic'
+
+  def get_flops(self, channels, check_range=True, divide=1):
+    iC, oC = channels
+    if check_range: assert iC <= self.conv.in_channels and oC <= self.conv.out_channels, '{:} vs {:}  |  {:} vs {:}'.format(iC, self.conv.in_channels, oC, self.conv.out_channels)
+    assert isinstance(self.InShape, tuple) and len(self.InShape) == 2, 'invalid in-shape : {:}'.format(self.InShape)
+    assert isinstance(self.OutShape, tuple) and len(self.OutShape) == 2, 'invalid out-shape : {:}'.format(self.OutShape)
+    #conv_per_position_flops = self.conv.kernel_size[0] * self.conv.kernel_size[1] * iC * oC / self.conv.groups
+    conv_per_position_flops = (self.conv.kernel_size[0] * self.conv.kernel_size[1] * 1.0 / self.conv.groups)
+    all_positions = self.OutShape[0] * self.OutShape[1]
+    flops = (conv_per_position_flops * all_positions / divide) * iC * oC
+    if self.conv.bias is not None: flops += all_positions / divide
+    return flops
+
+  def get_range(self):
+    return [self.choices]
+
+  def forward(self, inputs):
+    if self.search_mode == 'basic':
+      return self.basic_forward(inputs)
+    elif self.search_mode == 'search':
+      return self.search_forward(inputs)
+    else:
+      raise ValueError('invalid search_mode = {:}'.format(self.search_mode))
+
+  def search_forward(self, tuple_inputs):
+    assert isinstance(tuple_inputs, tuple) and len(tuple_inputs) == 5, 'invalid type input : {:}'.format( type(tuple_inputs) )
+    inputs, expected_inC, probability, index, prob = tuple_inputs
+    index, prob = torch.squeeze(index).tolist(), torch.squeeze(prob)
+    probability = torch.squeeze(probability)
+    assert len(index) == 2, 'invalid length : {:}'.format(index)
+    # compute expected flop
+    #coordinates   = torch.arange(self.x_range[0], self.x_range[1]+1).type_as(probability)
+    expected_outC = (self.choices_tensor * probability).sum()
+    expected_flop = self.get_flops([expected_inC, expected_outC], False, 1e6)
+    if self.avg : out = self.avg( inputs )
+    else        : out = inputs
+    # convolutional layer
+    out_convs = conv_forward(out, self.conv, [self.choices[i] for i in index])
+    out_bns   = [self.BNs[idx](out_conv) for idx, out_conv in zip(index, out_convs)]
+    # merge
+    out_channel = max([x.size(1) for x in out_bns])
+    outA = ChannelWiseInter(out_bns[0], out_channel)
+    outB = ChannelWiseInter(out_bns[1], out_channel)
+    out  = outA * prob[0] + outB * prob[1]
+    #out = additive_func(out_bns[0]*prob[0], out_bns[1]*prob[1])
+
+    if self.relu   : out = self.relu( out )
+    if self.maxpool: out = self.maxpool(out)
+    return out, expected_outC, expected_flop
+
+  def basic_forward(self, inputs):
+    if self.avg : out = self.avg( inputs )
+    else        : out = inputs
+    conv = self.conv( out )
+    if self.has_bn:out= self.BNs[-1]( conv )
+    else        : out = conv
+    if self.relu: out = self.relu( out )
+    else        : out = out
+    if self.InShape is None:
+      self.InShape  = (inputs.size(-2), inputs.size(-1))
+      self.OutShape = (out.size(-2)   , out.size(-1))
+    if self.maxpool: out = self.maxpool(out)
+    return out
+
+
+class ResNetBasicblock(nn.Module):
+  expansion = 1
+  num_conv  = 2
+  def __init__(self, inplanes, planes, stride):
+    super(ResNetBasicblock, self).__init__()
+    assert stride == 1 or stride == 2, 'invalid stride {:}'.format(stride)
+    self.conv_a = ConvBNReLU(inplanes, planes, 3, stride, 1, False, has_avg=False, has_bn=True, has_relu=True)
+    self.conv_b = ConvBNReLU(  planes, planes, 3,      1, 1, False, has_avg=False, has_bn=True, has_relu=False)
+    if stride == 2:
+      self.downsample = ConvBNReLU(inplanes, planes, 1, 1, 0, False, has_avg=True, has_bn=True, has_relu=False)
+    elif inplanes != planes:
+      self.downsample = ConvBNReLU(inplanes, planes, 1, 1, 0, False, has_avg=False,has_bn=True, has_relu=False)
+    else:
+      self.downsample = None
+    self.out_dim     = planes
+    self.search_mode = 'basic'
+
+  def get_range(self):
+    return self.conv_a.get_range() + self.conv_b.get_range()
+
+  def get_flops(self, channels):
+    assert len(channels) == 3, 'invalid channels : {:}'.format(channels)
+    flop_A = self.conv_a.get_flops([channels[0], channels[1]])
+    flop_B = self.conv_b.get_flops([channels[1], channels[2]])
+    if hasattr(self.downsample, 'get_flops'):
+      flop_C = self.downsample.get_flops([channels[0], channels[-1]])
+    else:
+      flop_C = 0
+    if channels[0] != channels[-1] and self.downsample is None: # this short-cut will be added during the infer-train
+      flop_C = channels[0] * channels[-1] * self.conv_b.OutShape[0] * self.conv_b.OutShape[1]
+    return flop_A + flop_B + flop_C
+
+  def forward(self, inputs):
+    if self.search_mode == 'basic'   : return self.basic_forward(inputs)
+    elif self.search_mode == 'search': return self.search_forward(inputs)
+    else: raise ValueError('invalid search_mode = {:}'.format(self.search_mode))
+
+  def search_forward(self, tuple_inputs):
+    assert isinstance(tuple_inputs, tuple) and len(tuple_inputs) == 5, 'invalid type input : {:}'.format( type(tuple_inputs) )
+    inputs, expected_inC, probability, indexes, probs = tuple_inputs
+    assert indexes.size(0) == 2 and probs.size(0) == 2 and probability.size(0) == 2
+    #import pdb; pdb.set_trace()
+    out_a, expected_inC_a, expected_flop_a = self.conv_a( (inputs, expected_inC  , probability[0], indexes[0], probs[0]) )
+    out_b, expected_inC_b, expected_flop_b = self.conv_b( (out_a , expected_inC_a, probability[1], indexes[1], probs[1]) )
+    if self.downsample is not None:
+      residual, _, expected_flop_c = self.downsample( (inputs, expected_inC  , probability[1], indexes[1], probs[1]) )
+    else:
+      residual, expected_flop_c = inputs, 0
+    out = additive_func(residual, out_b)
+    return out, expected_inC_b, sum([expected_flop_a, expected_flop_b, expected_flop_c])
+
+  def basic_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
+    out = additive_func(residual, basicblock)
+    return nn.functional.relu(out, inplace=True)
+
+
+
+class ResNetBottleneck(nn.Module):
+  expansion = 4
+  num_conv  = 3
+  def __init__(self, inplanes, planes, stride):
+    super(ResNetBottleneck, self).__init__()
+    assert stride == 1 or stride == 2, 'invalid stride {:}'.format(stride)
+    self.conv_1x1 = ConvBNReLU(inplanes, planes, 1,      1, 0, False, has_avg=False, has_bn=True, has_relu=True)
+    self.conv_3x3 = ConvBNReLU(  planes, planes, 3, stride, 1, False, has_avg=False, has_bn=True, has_relu=True)
+    self.conv_1x4 = ConvBNReLU(planes, planes*self.expansion, 1, 1, 0, False, has_avg=False, has_bn=True, has_relu=False)
+    if stride == 2:
+      self.downsample = ConvBNReLU(inplanes, planes*self.expansion, 1, 1, 0, False, has_avg=True, has_bn=True, has_relu=False)
+    elif inplanes != planes*self.expansion:
+      self.downsample = ConvBNReLU(inplanes, planes*self.expansion, 1, 1, 0, False, has_avg=False,has_bn=True, has_relu=False)
+    else:
+      self.downsample = None
+    self.out_dim     = planes * self.expansion
+    self.search_mode = 'basic'
+
+  def get_range(self):
+    return self.conv_1x1.get_range() + self.conv_3x3.get_range() + self.conv_1x4.get_range()
+
+  def get_flops(self, channels):
+    assert len(channels) == 4, 'invalid channels : {:}'.format(channels)
+    flop_A = self.conv_1x1.get_flops([channels[0], channels[1]])
+    flop_B = self.conv_3x3.get_flops([channels[1], channels[2]])
+    flop_C = self.conv_1x4.get_flops([channels[2], channels[3]])
+    if hasattr(self.downsample, 'get_flops'):
+      flop_D = self.downsample.get_flops([channels[0], channels[-1]])
+    else:
+      flop_D = 0
+    if channels[0] != channels[-1] and self.downsample is None: # this short-cut will be added during the infer-train
+      flop_D = channels[0] * channels[-1] * self.conv_1x4.OutShape[0] * self.conv_1x4.OutShape[1]
+    return flop_A + flop_B + flop_C + flop_D
+
+  def forward(self, inputs):
+    if self.search_mode == 'basic'   : return self.basic_forward(inputs)
+    elif self.search_mode == 'search': return self.search_forward(inputs)
+    else: raise ValueError('invalid search_mode = {:}'.format(self.search_mode))
+
+  def basic_forward(self, inputs):
+    bottleneck = self.conv_1x1(inputs)
+    bottleneck = self.conv_3x3(bottleneck)
+    bottleneck = self.conv_1x4(bottleneck)
+    if self.downsample is not None: residual = self.downsample(inputs)
+    else                          : residual = inputs
+    out = additive_func(residual, bottleneck)
+    return nn.functional.relu(out, inplace=True)
+
+  def search_forward(self, tuple_inputs):
+    assert isinstance(tuple_inputs, tuple) and len(tuple_inputs) == 5, 'invalid type input : {:}'.format( type(tuple_inputs) )
+    inputs, expected_inC, probability, indexes, probs = tuple_inputs
+    assert indexes.size(0) == 3 and probs.size(0) == 3 and probability.size(0) == 3
+    out_1x1, expected_inC_1x1, expected_flop_1x1 = self.conv_1x1( (inputs, expected_inC    , probability[0], indexes[0], probs[0]) )
+    out_3x3, expected_inC_3x3, expected_flop_3x3 = self.conv_3x3( (out_1x1,expected_inC_1x1, probability[1], indexes[1], probs[1]) )
+    out_1x4, expected_inC_1x4, expected_flop_1x4 = self.conv_1x4( (out_3x3,expected_inC_3x3, probability[2], indexes[2], probs[2]) )
+    if self.downsample is not None:
+      residual, _, expected_flop_c = self.downsample( (inputs, expected_inC  , probability[2], indexes[2], probs[2]) )
+    else:
+      residual, expected_flop_c = inputs, 0
+    out = additive_func(residual, out_1x4)
+    return out, expected_inC_1x4, sum([expected_flop_1x1, expected_flop_3x3, expected_flop_1x4, expected_flop_c])
+
+
+
+class SearchShapeImagenetResNet(nn.Module):
+
+  def __init__(self, block_name, layers, deep_stem, num_classes):
+    super(SearchShapeImagenetResNet, self).__init__()
+
+    #Model type specifies number of layers for CIFAR-10 and CIFAR-100 model
+    if block_name == 'BasicBlock':
+      block = ResNetBasicblock
+    elif block_name == 'Bottleneck':
+      block = ResNetBottleneck
+    else:
+      raise ValueError('invalid block : {:}'.format(block_name))
+    
+    self.message      = 'SearchShapeCifarResNet : Depth : {:} , Layers for each block : {:}'.format(sum(layers)*block.num_conv, layers)
+    self.num_classes  = num_classes
+    if not deep_stem:
+      self.layers       = nn.ModuleList( [ ConvBNReLU(3, 64, 7, 2, 3, False, has_avg=False, has_bn=True, has_relu=True, last_max_pool=True) ] )
+      self.channels     = [64]
+    else:
+      self.layers       = nn.ModuleList( [ ConvBNReLU(3, 32, 3, 2, 1, False, has_avg=False, has_bn=True, has_relu=True)
+                                          ,ConvBNReLU(32,64, 3, 1, 1, False, has_avg=False, has_bn=True, has_relu=True, last_max_pool=True) ] )
+      self.channels     = [32, 64]
+
+    meta_depth_info   = get_depth_choices(layers)
+    self.InShape      = None
+    self.depth_info   = OrderedDict()
+    self.depth_at_i   = OrderedDict()
+    for stage, layer_blocks in enumerate(layers):
+      cur_block_choices = meta_depth_info[stage]
+      assert cur_block_choices[-1] == layer_blocks, 'stage={:}, {:} vs {:}'.format(stage, cur_block_choices, layer_blocks)
+      block_choices, xstart = [], len(self.layers)
+      for iL in range(layer_blocks):
+        iC     = self.channels[-1]
+        planes = 64 * (2**stage)
+        stride = 2 if stage > 0 and iL == 0 else 1
+        module = block(iC, planes, stride)
+        self.channels.append( module.out_dim )
+        self.layers.append  ( module )
+        self.message += "\nstage={:}, ilayer={:02d}/{:02d}, block={:03d}, iC={:3d}, oC={:3d}, stride={:}".format(stage, iL, layer_blocks, len(self.layers)-1, iC, module.out_dim, stride)
+        # added for depth
+        layer_index = len(self.layers) - 1
+        if iL + 1 in cur_block_choices: block_choices.append( layer_index )
+        if iL + 1 == layer_blocks:
+          self.depth_info[layer_index] = {'choices': block_choices,
+                                          'stage'  : stage,
+                                          'xstart' : xstart}
+    self.depth_info_list = []
+    for xend, info in self.depth_info.items():
+      self.depth_info_list.append( (xend, info) )
+      xstart, xstage = info['xstart'], info['stage']
+      for ilayer in range(xstart, xend+1):
+        idx = bisect_right(info['choices'], ilayer-1)
+        self.depth_at_i[ilayer] = (xstage, idx)
+
+    self.avgpool     = nn.AdaptiveAvgPool2d((1,1))
+    self.classifier  = nn.Linear(module.out_dim, num_classes)
+    self.InShape     = None
+    self.tau         = -1
+    self.search_mode = 'basic'
+    #assert sum(x.num_conv for x in self.layers) + 1 == depth, 'invalid depth check {:} vs {:}'.format(sum(x.num_conv for x in self.layers)+1, depth)
+    
+    # parameters for width
+    self.Ranges = []
+    self.layer2indexRange = []
+    for i, layer in enumerate(self.layers):
+      start_index = len(self.Ranges)
+      self.Ranges += layer.get_range()
+      self.layer2indexRange.append( (start_index, len(self.Ranges)) )
+
+    self.register_parameter('width_attentions', nn.Parameter(torch.Tensor(len(self.Ranges), get_width_choices(None))))
+    self.register_parameter('depth_attentions', nn.Parameter(torch.Tensor(len(layers), meta_depth_info['num'])))
+    nn.init.normal_(self.width_attentions, 0, 0.01)
+    nn.init.normal_(self.depth_attentions, 0, 0.01)
+    self.apply(initialize_resnet)
+
+  def arch_parameters(self, LR=None):
+    if LR is None:
+      return [self.width_attentions, self.depth_attentions]
+    else:
+      return [
+               {"params": self.width_attentions, "lr": LR},
+               {"params": self.depth_attentions, "lr": LR},
+             ]
+
+  def base_parameters(self):
+    return list(self.layers.parameters()) + list(self.avgpool.parameters()) + list(self.classifier.parameters())
+
+  def get_flop(self, mode, config_dict, extra_info):
+    if config_dict is not None: config_dict = config_dict.copy()
+    # select channels 
+    channels = [3]
+    for i, weight in enumerate(self.width_attentions):
+      if mode == 'genotype':
+        with torch.no_grad():
+          probe = nn.functional.softmax(weight, dim=0)
+          C = self.Ranges[i][ torch.argmax(probe).item() ]
+      else:
+        raise ValueError('invalid mode : {:}'.format(mode))
+      channels.append( C )
+    # select depth
+    if mode == 'genotype':
+      with torch.no_grad():
+        depth_probs = nn.functional.softmax(self.depth_attentions, dim=1)
+        choices = torch.argmax(depth_probs, dim=1).cpu().tolist()
+    else:
+      raise ValueError('invalid mode : {:}'.format(mode))
+    selected_layers = []
+    for choice, xvalue in zip(choices, self.depth_info_list):
+      xtemp = xvalue[1]['choices'][choice] - xvalue[1]['xstart'] + 1
+      selected_layers.append(xtemp)
+    flop = 0
+    for i, layer in enumerate(self.layers):
+      s, e = self.layer2indexRange[i]
+      xchl = tuple( channels[s:e+1] )
+      if i in self.depth_at_i:
+        xstagei, xatti = self.depth_at_i[i]
+        if xatti <= choices[xstagei]: # leave this depth
+          flop+= layer.get_flops(xchl)
+        else:
+          flop+= 0 # do not use this layer
+      else:
+        flop+= layer.get_flops(xchl)
+    # the last fc layer
+    flop += channels[-1] * self.classifier.out_features
+    if config_dict is None:
+      return flop / 1e6
+    else:
+      config_dict['xchannels']  = channels
+      config_dict['xblocks']    = selected_layers
+      config_dict['super_type'] = 'infer-shape'
+      config_dict['estimated_FLOP'] = flop / 1e6
+      return flop / 1e6, config_dict
+
+  def get_arch_info(self):
+    string = "for depth and width, there are {:} + {:} attention probabilities.".format(len(self.depth_attentions), len(self.width_attentions))
+    string+= '\n{:}'.format(self.depth_info)
+    discrepancy = []
+    with torch.no_grad():
+      for i, att in enumerate(self.depth_attentions):
+        prob = nn.functional.softmax(att, dim=0)
+        prob = prob.cpu() ; selc = prob.argmax().item() ; prob = prob.tolist()
+        prob = ['{:.3f}'.format(x) for x in prob]
+        xstring = '{:03d}/{:03d}-th : {:}'.format(i, len(self.depth_attentions), ' '.join(prob))
+        logt = ['{:.4f}'.format(x) for x in att.cpu().tolist()]
+        xstring += '  ||  {:17s}'.format(' '.join(logt))
+        prob = sorted( [float(x) for x in prob] )
+        disc = prob[-1] - prob[-2]
+        xstring += '  || discrepancy={:.2f} || select={:}/{:}'.format(disc, selc, len(prob))
+        discrepancy.append( disc )
+        string += '\n{:}'.format(xstring)
+      string += '\n-----------------------------------------------'
+      for i, att in enumerate(self.width_attentions):
+        prob = nn.functional.softmax(att, dim=0)
+        prob = prob.cpu() ; selc = prob.argmax().item() ; prob = prob.tolist()
+        prob = ['{:.3f}'.format(x) for x in prob]
+        xstring = '{:03d}/{:03d}-th : {:}'.format(i, len(self.width_attentions), ' '.join(prob))
+        logt = ['{:.3f}'.format(x) for x in att.cpu().tolist()]
+        xstring += '  ||  {:52s}'.format(' '.join(logt))
+        prob = sorted( [float(x) for x in prob] )
+        disc = prob[-1] - prob[-2]
+        xstring += '  || dis={:.2f} || select={:}/{:}'.format(disc, selc, len(prob))
+        discrepancy.append( disc )
+        string += '\n{:}'.format(xstring)
+    return string, discrepancy
+
+  def set_tau(self, tau_max, tau_min, epoch_ratio):
+    assert epoch_ratio >= 0 and epoch_ratio <= 1, 'invalid epoch-ratio : {:}'.format(epoch_ratio)
+    tau = tau_min + (tau_max-tau_min) * (1 + math.cos(math.pi * epoch_ratio)) / 2
+    self.tau = tau
+
+  def get_message(self):
+    return self.message
+
+  def forward(self, inputs):
+    if self.search_mode == 'basic':
+      return self.basic_forward(inputs)
+    elif self.search_mode == 'search':
+      return self.search_forward(inputs)
+    else:
+      raise ValueError('invalid search_mode = {:}'.format(self.search_mode))
+
+  def search_forward(self, inputs):
+    flop_width_probs = nn.functional.softmax(self.width_attentions, dim=1)
+    flop_depth_probs = nn.functional.softmax(self.depth_attentions, dim=1)
+    flop_depth_probs = torch.flip( torch.cumsum( torch.flip(flop_depth_probs, [1]), 1 ), [1] )
+    selected_widths, selected_width_probs = select2withP(self.width_attentions, self.tau)
+    selected_depth_probs = select2withP(self.depth_attentions, self.tau, True)
+    with torch.no_grad():
+      selected_widths = selected_widths.cpu()
+
+    x, last_channel_idx, expected_inC, flops = inputs, 0, 3, []
+    feature_maps = []
+    for i, layer in enumerate(self.layers):
+      selected_w_index = selected_widths     [last_channel_idx: last_channel_idx+layer.num_conv]
+      selected_w_probs = selected_width_probs[last_channel_idx: last_channel_idx+layer.num_conv]
+      layer_prob       = flop_width_probs    [last_channel_idx: last_channel_idx+layer.num_conv]
+      x, expected_inC, expected_flop = layer( (x, expected_inC, layer_prob, selected_w_index, selected_w_probs) )
+      feature_maps.append( x )
+      last_channel_idx += layer.num_conv
+      if i in self.depth_info: # aggregate the information
+        choices = self.depth_info[i]['choices']
+        xstagei = self.depth_info[i]['stage']
+        #print ('iL={:}, choices={:}, stage={:}, probs={:}'.format(i, choices, xstagei, selected_depth_probs[xstagei].cpu().tolist()))
+        #for A, W in zip(choices, selected_depth_probs[xstagei]):
+        #  print('Size = {:}, W = {:}'.format(feature_maps[A].size(), W))
+        possible_tensors = []
+        max_C = max( feature_maps[A].size(1) for A in choices )
+        for tempi, A in enumerate(choices):
+          xtensor = ChannelWiseInter(feature_maps[A], max_C)
+          possible_tensors.append( xtensor )
+        weighted_sum = sum( xtensor * W for xtensor, W in zip(possible_tensors, selected_depth_probs[xstagei]) )
+        x = weighted_sum
+        
+      if i in self.depth_at_i:
+        xstagei, xatti = self.depth_at_i[i]
+        x_expected_flop = flop_depth_probs[xstagei, xatti] * expected_flop
+      else:
+        x_expected_flop = expected_flop
+      flops.append( x_expected_flop )
+    flops.append( expected_inC * (self.classifier.out_features*1.0/1e6) )
+    features = self.avgpool(x)
+    features = features.view(features.size(0), -1)
+    logits   = linear_forward(features, self.classifier)
+    return logits, torch.stack( [sum(flops)] )
+
+  def basic_forward(self, inputs):
+    if self.InShape is None: self.InShape = (inputs.size(-2), inputs.size(-1))
+    x = inputs
+    for i, layer in enumerate(self.layers):
+      x = layer( x )
+    features = self.avgpool(x)
+    features = features.view(features.size(0), -1)
+    logits   = self.classifier(features)
+    return features, logits

lib/models/searchs/SoftSelect.py

@@ -0,0 +1,108 @@
+import math, torch
+import torch.nn as nn
+
+
+def select2withP(logits, tau, just_prob=False, num=2, eps=1e-7):
+  if tau <= 0:
+    new_logits = logits
+    probs = nn.functional.softmax(new_logits, dim=1)
+  else       :
+    while True: # a trick to avoid the gumbels bug
+      gumbels = -torch.empty_like(logits).exponential_().log()
+      new_logits = (logits + gumbels) / tau
+      probs = nn.functional.softmax(new_logits, dim=1)
+      if (not torch.isinf(gumbels).any()) and (not torch.isinf(probs).any()) and (not torch.isnan(probs).any()): break
+
+  if just_prob: return probs
+
+  #with torch.no_grad(): # add eps for unexpected torch error
+  #  probs = nn.functional.softmax(new_logits, dim=1)
+  #  selected_index = torch.multinomial(probs + eps, 2, False)
+  with torch.no_grad(): # add eps for unexpected torch error
+    probs          = probs.cpu()
+    selected_index = torch.multinomial(probs + eps, num, False).to(logits.device)
+  selected_logit = torch.gather(new_logits, 1, selected_index)
+  selcted_probs  = nn.functional.softmax(selected_logit, dim=1)
+  return selected_index, selcted_probs
+
+
+def ChannelWiseInter(inputs, oC, mode='v2'):
+  if mode == 'v1':
+    return ChannelWiseInterV1(inputs, oC)
+  elif mode == 'v2':
+    return ChannelWiseInterV2(inputs, oC)
+  else:
+    raise ValueError('invalid mode : {:}'.format(mode))
+
+
+def ChannelWiseInterV1(inputs, oC):
+  assert inputs.dim() == 4, 'invalid dimension : {:}'.format(inputs.size())
+  def start_index(a, b, c):
+    return int( math.floor(float(a * c) / b) )
+  def end_index(a, b, c):
+    return int( math.ceil(float((a + 1) * c) / b) )
+  batch, iC, H, W = inputs.size()
+  outputs = torch.zeros((batch, oC, H, W), dtype=inputs.dtype, device=inputs.device)
+  if iC == oC: return inputs
+  for ot in range(oC):
+    istartT, iendT = start_index(ot, oC, iC), end_index(ot, oC, iC)
+    values = inputs[:, istartT:iendT].mean(dim=1) 
+    outputs[:, ot, :, :] = values
+  return outputs
+
+
+def ChannelWiseInterV2(inputs, oC):
+  assert inputs.dim() == 4, 'invalid dimension : {:}'.format(inputs.size())
+  batch, C, H, W = inputs.size()
+  if C == oC: return inputs
+  else      : return nn.functional.adaptive_avg_pool3d(inputs, (oC,H,W))
+  #inputs_5D = inputs.view(batch, 1, C, H, W)
+  #otputs_5D = nn.functional.interpolate(inputs_5D, (oC,H,W), None, 'area', None)
+  #otputs    = otputs_5D.view(batch, oC, H, W)
+  #otputs_5D = nn.functional.interpolate(inputs_5D, (oC,H,W), None, 'trilinear', False)
+  #return otputs
+
+
+def linear_forward(inputs, linear):
+  if linear is None: return inputs
+  iC = inputs.size(1)
+  weight = linear.weight[:, :iC]
+  if linear.bias is None: bias = None
+  else                  : bias = linear.bias
+  return nn.functional.linear(inputs, weight, bias)
+
+
+def get_width_choices(nOut):
+  xsrange = [0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
+  if nOut is None:
+    return len(xsrange)
+  else:
+    Xs = [int(nOut * i) for i in xsrange]
+    #xs = [ int(nOut * i // 10) for i in range(2, 11)]
+    #Xs = [x for i, x in enumerate(xs) if i+1 == len(xs) or xs[i+1] > x+1]
+    Xs = sorted( list( set(Xs) ) )
+    return tuple(Xs)
+
+
+def get_depth_choices(nDepth):
+  if nDepth is None:
+    return 3
+  else:
+    assert nDepth >= 3, 'nDepth should be greater than 2 vs {:}'.format(nDepth)
+    if nDepth == 1  : return (1, 1, 1)
+    elif nDepth == 2: return (1, 1, 2)
+    elif nDepth >= 3:
+      return (nDepth//3, nDepth*2//3, nDepth)
+    else:
+      raise ValueError('invalid Depth : {:}'.format(nDepth))
+
+
+def drop_path(x, drop_prob):
+  if drop_prob > 0.:
+    keep_prob = 1. - drop_prob
+    mask = x.new_zeros(x.size(0), 1, 1, 1)
+    mask = mask.bernoulli_(keep_prob)
+    x = x * (mask / keep_prob)
+    #x.div_(keep_prob)
+    #x.mul_(mask)
+  return x

lib/models/searchs/__init__.py

@@ -0,0 +1,4 @@
+from .SearchCifarResNet_width import SearchWidthCifarResNet
+from .SearchCifarResNet_depth import SearchDepthCifarResNet
+from .SearchCifarResNet       import SearchShapeCifarResNet
+from .SearchImagenetResNet    import SearchShapeImagenetResNet

lib/models/searchs/test.py

@@ -0,0 +1,17 @@
+import torch
+import torch.nn as nn
+from SoftSelect import ChannelWiseInter
+
+
+if __name__ == '__main__':
+
+  tensors = torch.rand((16, 128, 7, 7))
+  
+  for oc in range(200, 210):
+    out_v1  = ChannelWiseInter(tensors, oc, 'v1')
+    out_v2  = ChannelWiseInter(tensors, oc, 'v2')
+    assert (out_v1 == out_v2).any().item() == 1
+  for oc in range(48, 160):
+    out_v1  = ChannelWiseInter(tensors, oc, 'v1')
+    out_v2  = ChannelWiseInter(tensors, oc, 'v2')
+    assert (out_v1 == out_v2).any().item() == 1

lib/models/sphereface.py

@@ -0,0 +1,139 @@
+# SphereFace: Deep Hypersphere Embedding for Face Recognition
+#
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+def myphi(x,m):
+  x = x * m
+  return 1-x**2/math.factorial(2)+x**4/math.factorial(4)-x**6/math.factorial(6) + \
+      x**8/math.factorial(8) - x**9/math.factorial(9)
+
+class AngleLinear(nn.Module):
+  def __init__(self, in_features, out_features, m = 4, phiflag=True):
+    super(AngleLinear, self).__init__()
+    self.in_features = in_features
+    self.out_features = out_features
+    self.weight = nn.Parameter(torch.Tensor(in_features,out_features))
+    self.weight.data.uniform_(-1, 1).renorm_(2,1,1e-5).mul_(1e5)
+    self.phiflag = phiflag
+    self.m = m
+    self.mlambda = [
+      lambda x: x**0,
+      lambda x: x**1,
+      lambda x: 2*x**2-1,
+      lambda x: 4*x**3-3*x,
+      lambda x: 8*x**4-8*x**2+1,
+      lambda x: 16*x**5-20*x**3+5*x
+    ]
+
+  def forward(self, input):
+    x = input   # size=(B,F)  F is feature len
+    w = self.weight # size=(F,Classnum) F=in_features Classnum=out_features
+
+    ww = w.renorm(2,1,1e-5).mul(1e5)
+    xlen = x.pow(2).sum(1).pow(0.5) # size=B
+    wlen = ww.pow(2).sum(0).pow(0.5) # size=Classnum
+
+    cos_theta = x.mm(ww) # size=(B,Classnum)
+    cos_theta = cos_theta / xlen.view(-1,1) / wlen.view(1,-1)
+    cos_theta = cos_theta.clamp(-1,1)
+
+    if self.phiflag:
+      cos_m_theta = self.mlambda[self.m](cos_theta)
+      with torch.no_grad():
+        theta = cos_theta.acos()
+      k = (self.m*theta/3.14159265).floor()
+      n_one = k*0.0 - 1
+      phi_theta = (n_one**k) * cos_m_theta - 2*k
+    else:
+      theta = cos_theta.acos()
+      phi_theta = myphi(theta,self.m)
+      phi_theta = phi_theta.clamp(-1*self.m,1)
+
+    cos_theta = cos_theta * xlen.view(-1,1)
+    phi_theta = phi_theta * xlen.view(-1,1)
+    output = (cos_theta,phi_theta)
+    return output # size=(B,Classnum,2)
+
+
+class SphereFace20(nn.Module):
+  def __init__(self, classnum=10574):
+    super(SphereFace20, self).__init__()
+    self.classnum = classnum
+    #input = B*3*112*96
+    self.conv1_1 = nn.Conv2d(3,64,3,2,1) #=>B*64*56*48
+    self.relu1_1 = nn.PReLU(64)
+    self.conv1_2 = nn.Conv2d(64,64,3,1,1)
+    self.relu1_2 = nn.PReLU(64)
+    self.conv1_3 = nn.Conv2d(64,64,3,1,1)
+    self.relu1_3 = nn.PReLU(64)
+
+    self.conv2_1 = nn.Conv2d(64,128,3,2,1) #=>B*128*28*24
+    self.relu2_1 = nn.PReLU(128)
+    self.conv2_2 = nn.Conv2d(128,128,3,1,1)
+    self.relu2_2 = nn.PReLU(128)
+    self.conv2_3 = nn.Conv2d(128,128,3,1,1)
+    self.relu2_3 = nn.PReLU(128)
+
+    self.conv2_4 = nn.Conv2d(128,128,3,1,1) #=>B*128*28*24
+    self.relu2_4 = nn.PReLU(128)
+    self.conv2_5 = nn.Conv2d(128,128,3,1,1)
+    self.relu2_5 = nn.PReLU(128)
+
+
+    self.conv3_1 = nn.Conv2d(128,256,3,2,1) #=>B*256*14*12
+    self.relu3_1 = nn.PReLU(256)
+    self.conv3_2 = nn.Conv2d(256,256,3,1,1)
+    self.relu3_2 = nn.PReLU(256)
+    self.conv3_3 = nn.Conv2d(256,256,3,1,1)
+    self.relu3_3 = nn.PReLU(256)
+
+    self.conv3_4 = nn.Conv2d(256,256,3,1,1) #=>B*256*14*12
+    self.relu3_4 = nn.PReLU(256)
+    self.conv3_5 = nn.Conv2d(256,256,3,1,1)
+    self.relu3_5 = nn.PReLU(256)
+
+    self.conv3_6 = nn.Conv2d(256,256,3,1,1) #=>B*256*14*12
+    self.relu3_6 = nn.PReLU(256)
+    self.conv3_7 = nn.Conv2d(256,256,3,1,1)
+    self.relu3_7 = nn.PReLU(256)
+
+    self.conv3_8 = nn.Conv2d(256,256,3,1,1) #=>B*256*14*12
+    self.relu3_8 = nn.PReLU(256)
+    self.conv3_9 = nn.Conv2d(256,256,3,1,1)
+    self.relu3_9 = nn.PReLU(256)
+
+    self.conv4_1 = nn.Conv2d(256,512,3,2,1) #=>B*512*7*6
+    self.relu4_1 = nn.PReLU(512)
+    self.conv4_2 = nn.Conv2d(512,512,3,1,1)
+    self.relu4_2 = nn.PReLU(512)
+    self.conv4_3 = nn.Conv2d(512,512,3,1,1)
+    self.relu4_3 = nn.PReLU(512)
+
+    self.fc5 = nn.Linear(512*7*6,512)
+    self.fc6 = AngleLinear(512, self.classnum)
+
+
+  def forward(self, x):
+    x = self.relu1_1(self.conv1_1(x))
+    x = x + self.relu1_3(self.conv1_3(self.relu1_2(self.conv1_2(x))))
+
+    x = self.relu2_1(self.conv2_1(x))
+    x = x + self.relu2_3(self.conv2_3(self.relu2_2(self.conv2_2(x))))
+    x = x + self.relu2_5(self.conv2_5(self.relu2_4(self.conv2_4(x))))
+
+    x = self.relu3_1(self.conv3_1(x))
+    x = x + self.relu3_3(self.conv3_3(self.relu3_2(self.conv3_2(x))))
+    x = x + self.relu3_5(self.conv3_5(self.relu3_4(self.conv3_4(x))))
+    x = x + self.relu3_7(self.conv3_7(self.relu3_6(self.conv3_6(x))))
+    x = x + self.relu3_9(self.conv3_9(self.relu3_8(self.conv3_8(x))))
+
+    x = self.relu4_1(self.conv4_1(x))
+    x = x + self.relu4_3(self.conv4_3(self.relu4_2(self.conv4_2(x))))
+
+    x = x.view(x.size(0),-1)
+    features = self.fc5(x)
+    logits   = self.fc6(features)
+    return features, logits