v0.1 data_parallel support (bn track_stats) + preResNet/denseNet/resNeXt

Author: solopku

models/__init__.py

@@ -1,2 +1,5 @@
 from .convBlock import *
-from .wideResnet import get_wrn
+from .wideResnet import get_wrn
+from .preResNet import get_resnet
+from .denseNet import get_densenet
+from .resNeXt import get_resnext

models/convBlock.py

@@ -6,18 +6,21 @@ def __init__(self, drop_type=0, drop_rate=0., inplace=False):
         super(drop_op, self).__init__()
         assert (drop_type in (0, 1, 2, 3) and 0.<=drop_rate<1.)
         self.drop_type, self.keep_rate = drop_type, 1.-drop_rate
+        if drop_rate == 0.:
+            self.drop_op = nn.Sequential();return
         if drop_type == 0:
             self.drop_op = nn.Dropout(p=drop_rate, inplace=inplace)
         elif drop_type == 1:
             self.drop_op = nn.Dropout2d(p=drop_rate, inplace=inplace)
 
     def forward(self, x):
+        if self.keep_rate == 1.: return x
         if self.drop_type in (0, 1): return self.drop_op(x)
         # drop-branch/layer, x in [b_0, b_1, ...], b_i: B*C_i*H*W
         if self.training:
-            mask = torch.FloatTensor(1, x.size(1), 1, 1, device=x.device).\
+            mask = torch.FloatTensor(len(x)).to(x[0].device).\
                 bernoulli_(self.keep_rate)*(1./self.keep_rate)
-            x = list(map(lambda b: b.mul_(mask), x))
+            x = [x[idx]*mask[idx] for idx in range(len(x))]
         return torch.cat(x, dim=1)
 
 class Norm2d(nn.Module):
@@ -46,15 +49,11 @@ def forward(self, input):
                 var = input.var([0, 2, 3], unbiased=False)
                 n = input.numel()/input.size(1)
                 if self.training:
-                    self.train_running_mean = self.momentum*mean+\
-                                        (1-self.momentum)*self.train_running_mean
-                    self.train_running_var = self.momentum*var*n/(n-1)+\
-                                        (1-self.momentum)*self.train_running_var
+                    self.train_running_mean.mul_(1 - self.momentum).add_(self.momentum*mean)
+                    self.train_running_var.mul_(1-self.momentum).add_(self.momentum*var*n/(n-1))
                 else:
-                    self.test_running_mean = self.momentum*mean+\
-                                        (1-self.momentum)*self.test_running_mean
-                    self.test_running_var = self.momentum*var*n/(n-1)+\
-                                        (1-self.momentum)*self.test_running_var
+                    self.test_running_mean.mul_(1 - self.momentum).add_(self.momentum*mean)
+                    self.test_running_var.mul_(1-self.momentum).add_(self.momentum*var*n/(n-1))
         return self.norm(input)
 
 def norm2d_track_stats(model, is_track):
@@ -72,13 +71,14 @@ def norm2d_stats(model):
 class conv_block(nn.Module):
     def __init__(self, in_channels, out_channels, kernel_size, block_type=0,
                  use_gn=False, gn_groups=8,  drop_type=0, drop_rate=0.,
-                 stride=1, padding=0, groups=1, bias=False):
+                 stride=1, padding=0, groups=1, bias=False, track_stats=False):
         super(conv_block, self).__init__()
         self.relu = nn.ReLU(inplace=True)
-        self.norm = Norm2d(in_channels, use_gn, gn_groups, drop_rate>0.)
         self.drop = drop_op(drop_type, drop_rate)
         self.conv = nn.Conv2d(in_channels, out_channels, kernel_size,
                 groups=groups, stride=stride, padding=padding, bias=bias)
+        bn_channels = in_channels if block_type in [0, 1] else out_channels
+        self.norm = Norm2d(bn_channels, use_gn, gn_groups, track_stats)
 
         if block_type==0:       # bn/gn-relu-drop-conv, recommended
             self.ops = nn.Sequential(self.norm, self.relu,

models/denseNet.py

@@ -0,0 +1,132 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from models import conv_block
+
+class BasicBlock(nn.Module):
+    def __init__(self, in_planes, out_planes, args):
+        super(BasicBlock, self).__init__()
+        self.conv = conv_block(in_planes, out_planes, 3, args.block_type,
+                        args.use_gn, args.gn_groups, args.drop_type,
+                        args.drop_rate, padding=1, track_stats=args.report_ratio)
+
+    def forward(self, x):
+        out = self.conv(x)
+        return torch.cat([x, out], 1)
+
+class Bottleneck(nn.Module):
+    def __init__(self, in_planes, out_planes, args):
+        super(Bottleneck, self).__init__()
+        inter_planes = out_planes * 4
+        self.conv1 = conv_block(in_planes, inter_planes, 1, args.block_type,
+                                args.use_gn, args.gn_groups, args.drop_type, args.drop_rate,
+                                track_stats=args.report_ratio)
+        self.conv2 = conv_block(inter_planes, out_planes, 3, args.block_type,
+                                args.use_gn, args.gn_groups, args.drop_type, args.drop_rate,
+                                padding=1, track_stats=args.report_ratio)
+
+    def forward(self, x):
+        out = self.conv2(self.conv1(x))
+        return torch.cat([x, out], 1)
+
+class TransitionBlock(nn.Module):
+    def __init__(self, in_planes, out_planes, args):
+        super(TransitionBlock, self).__init__()
+        self.conv = conv_block(in_planes, out_planes, 1, args.block_type,
+                        args.use_gn, args.gn_groups, args.drop_type, args.drop_rate,
+                        track_stats=args.report_ratio)
+
+    def forward(self, x):
+        out = self.conv(x)
+        return F.avg_pool2d(out, 2)
+
+class DenseBlock(nn.Module):
+    def __init__(self, num_layers, in_planes, growth_rate, block, args):
+        super(DenseBlock, self).__init__()
+        self.layer = nn.Sequential(*[block(in_planes+i*growth_rate, growth_rate, args)
+                                     for i in range(num_layers)])
+
+    def forward(self, x):
+        return self.layer(x)
+
+# For CIFAR-10/100 dataset
+class DenseNet(nn.Module):
+    def __init__(self, args, growth_rate=12,
+                 reduction=0.5, bottleneck=True):
+        super(DenseNet, self).__init__()
+        in_planes = 2 * growth_rate
+        n = int((args.depth - 4) / 3)
+        if bottleneck == True:
+            n = n//2
+            block = Bottleneck
+        else:
+            block = BasicBlock
+        # 1st conv before any dense block
+        self.conv1 = nn.Conv2d(3, in_planes, kernel_size=3, stride=1,
+                               padding=1, bias=False)
+        # 1st block
+        self.block1 = DenseBlock(n, in_planes, growth_rate, block, args)
+        in_planes = int(in_planes+n*growth_rate)
+        self.trans1 = TransitionBlock(in_planes, int(math.floor(in_planes*reduction)), args)
+        in_planes = int(math.floor(in_planes*reduction))
+        # 2nd block
+        self.block2 = DenseBlock(n, in_planes, growth_rate, block, args)
+        in_planes = int(in_planes+n*growth_rate)
+        self.trans2 = TransitionBlock(in_planes, int(math.floor(in_planes*reduction)), args)
+        in_planes = int(math.floor(in_planes*reduction))
+        # 3rd block
+        self.block3 = DenseBlock(n, in_planes, growth_rate, block, args)
+        in_planes = int(in_planes+n*growth_rate)
+        # global average pooling and classifier
+        self.bn = nn.BatchNorm2d(in_planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.fc = nn.Linear(in_planes, args.class_num)
+        self.in_planes = in_planes
+
+    def forward(self, x):
+        out = self.conv1(x)
+        out = self.trans1(self.block1(out))
+        out = self.trans2(self.block2(out))
+        out = self.block3(out)
+        out = self.relu(self.bn(out))
+        out = F.avg_pool2d(out, 8)
+        out = out.view(-1, self.in_planes)
+        return self.fc(out)
+
+
+# https://github.com/liuzhuang13/DenseNet#results-on-cifar
+# CIFAR DenseNet3(depth=100, num_classes=10., growth_rate=12, reduction=0.5, bottleneck=True, drop_rate=0.2)
+# SVHN  DenseNet3(depth=100, num_classes=10., growth_rate=24, reduction=0.5, bottleneck=True, drop_rate=0.2)
+# DenseNet3(depth=250, num_classes=10., growth_rate=24, reduction=0.5, bottleneck=True, drop_rate=0.2)
+# DenseNet3(depth=190, num_classes=10., growth_rate=40, reduction=0.5, bottleneck=True, drop_rate=0.2)
+def get_densenet(args):
+    return DenseNet(args, args.arg1)
+
+if __name__ == '__main__':
+    import argparse
+
+    parser = argparse.ArgumentParser(description='WideResNet')
+    args = parser.parse_args()
+    args.depth = 100
+    args.class_num = 10
+    args.block_type = 0
+    args.use_gn = False
+    args.gn_groups = 6
+    args.drop_type = 1
+    args.drop_rate = 0.1
+    args.report_ratio = True
+    args.arg1 = 12
+
+    net = DenseNet(args, args.arg1)
+    y = net(torch.randn(1, 3, 32, 32))
+    print(y.size())
+    print(net)
+    print(sum([p.data.nelement() for p in net.parameters()]))
+
+    from convBlock import Norm2d, norm2d_stats, norm2d_track_stats
+
+    # norm2d_track_stats(net, False)
+    mean, var = norm2d_stats(net)
+    print(len(mean), mean)
+    print(var)

models/preResNet.py

@@ -0,0 +1,113 @@
+import torch
+import torch.nn as nn
+from models import conv_block
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, args, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+
+        self.bn1 = nn.BatchNorm2d(inplanes)
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.relu = nn.ReLU(inplace=True)
+
+        self.block2 = conv_block(planes, planes, 3, args.block_type, args.use_gn, args.gn_groups,
+                                 args.drop_type, args.drop_rate, stride=stride, padding=1,
+                                 track_stats=args.report_ratio)
+        self.block3 = conv_block(planes, planes*Bottleneck.expansion, 1, block_type=0,
+                                 use_gn=False, drop_rate=0., track_stats=False)
+
+        self.downsample = downsample
+
+    def forward(self, x):
+        residual = x
+
+        out = self.relu(self.bn1(x))
+
+        if self.downsample is not None:
+            residual = self.downsample(out)
+
+        out = self.conv1(out)
+        out = self.block2(out)
+        out = self.block3(out)
+
+        out += residual
+
+        return out
+
+
+class preResNet(nn.Module):
+    def __init__(self, args, widen_factor=1.):
+        super(preResNet, self).__init__()
+        self.inplanes = int(16*widen_factor)
+        n = int((args.depth - 2) / 9)
+
+        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.layer1 = self._make_layer(int(16*widen_factor), n, args)
+        self.layer2 = self._make_layer(int(32*widen_factor), n, args, stride=2)
+        self.layer3 = self._make_layer(int(64*widen_factor), n, args, stride=2)
+        self.bn = nn.BatchNorm2d(int(64 * widen_factor) * Bottleneck.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.avgpool = nn.AvgPool2d(8)
+        self.fc = nn.Linear(int(64*widen_factor) * Bottleneck.expansion, args.class_num)
+
+    def _make_layer(self, planes, blocks, args, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * Bottleneck.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * Bottleneck.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes * Bottleneck.expansion),
+            )
+
+        layers = []
+        layers.append(Bottleneck(self.inplanes, planes, args, stride, downsample))
+        self.inplanes = planes * Bottleneck.expansion
+        for i in range(1, blocks):
+            layers.append(Bottleneck(self.inplanes, planes, args))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+
+        x = self.relu(self.bn(x))
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+
+        return x
+
+def get_resnet(args):
+    return preResNet(args, args.arg1)
+
+if __name__ == '__main__':
+    import argparse
+    parser = argparse.ArgumentParser(description='PreResNet')
+    args = parser.parse_args()
+    args.depth=110
+    args.class_num = 10
+    args.block_type = 0
+    args.use_gn = False
+    args.gn_groups = 16
+    args.drop_type = 1
+    args.drop_rate = 0.1
+    args.report_ratio = True
+
+    net = preResNet(args)
+    y = net((torch.randn(1, 3, 32, 32)))
+    print(y.size())
+    print(net)
+    print(sum([p.data.nelement() for p in net.parameters()]))
+
+    from convBlock import Norm2d, norm2d_stats, norm2d_track_stats
+
+    # norm2d_track_stats(net, False)
+    mean, var = norm2d_stats(net)
+    print(mean)
+    print(var)