added upconv and upproj

Author: Fangchang Ma

main.py

@@ -53,8 +53,8 @@
                         ' (default: deconv2)')
 parser.add_argument('-j', '--workers', default=10, type=int, metavar='N',
                     help='number of data loading workers (default: 10)')
-parser.add_argument('--epochs', default=30, type=int, metavar='N',
-                    help='number of total epochs to run (default: 30)')
+parser.add_argument('--epochs', default=15, type=int, metavar='N',
+                    help='number of total epochs to run (default: 15)')
 parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
                     help='manual epoch number (useful on restarts)')
 parser.add_argument('-c', '--criterion', metavar='LOSS', default='l1', 
@@ -106,10 +106,9 @@ def main():
     # define loss function (criterion) and optimizer
     if args.criterion == 'l2':
         criterion = criteria.MaskedMSELoss().cuda()
-        out_channels = 1
     elif args.criterion == 'l1':
         criterion = criteria.MaskedL1Loss().cuda()
-        out_channels = 1
+    out_channels = 1
 
     # Data loading code
     print("=> creating data loaders ...")
@@ -157,6 +156,7 @@ def main():
             print("=> loaded checkpoint (epoch {})".format(checkpoint['epoch']))
         else:
             print("=> no checkpoint found at '{}'".format(args.resume))
+            return
 
     # create new model
     else:
@@ -358,8 +358,8 @@ def save_checkpoint(state, is_best, epoch):
             os.remove(prev_checkpoint_filename)
 
 def adjust_learning_rate(optimizer, epoch):
-    """Sets the learning rate to the initial LR decayed by 10 every 10 epochs"""
-    lr = args.lr * (0.1 ** (epoch // 10))
+    """Sets the learning rate to the initial LR decayed by 10 every 5 epochs"""
+    lr = args.lr * (0.1 ** (epoch // 5))
     for param_group in optimizer.param_groups:
         param_group['lr'] = lr
 

models.py

@@ -1,12 +1,28 @@
 import os
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 import torchvision.models
 import collections
 import math
 
 oheight, owidth = 228, 304
 
+class Unpool(nn.Module):
+    # Unpool: 2*2 unpooling with zero padding 
+    def __init__(self, num_channels, stride=2):
+        super(Unpool, self).__init__()
+
+        self.num_channels = num_channels
+        self.stride = stride
+
+        # create kernel [1, 0; 0, 0]
+        self.weights = torch.autograd.Variable(torch.zeros(num_channels, 1, stride, stride).cuda()) # currently not compatible with running on CPU 
+        self.weights[:,:,0,0] = 1
+
+    def forward(self, x):
+        return F.conv_transpose2d(x, self.weights, stride=self.stride, groups=self.num_channels)
+
 def weights_init(m):
     # Initialize filters with Gaussian random weights
     if isinstance(m, nn.Conv2d): 
@@ -26,7 +42,7 @@ def weights_init(m):
 class Decoder(nn.Module):
     # Decoder is the base class for all decoders
 
-    names = ['deconv{}'.format(i) for i in range(2,10)]
+    names = ['deconv2', 'deconv3', 'upconv', 'upproj']
 
     def __init__(self):
         super(Decoder, self).__init__()
@@ -67,15 +83,77 @@ def convt(in_channels):
         self.layer3 = convt(in_channels // (2 ** 2))
         self.layer4 = convt(in_channels // (2 ** 3))
 
-
-def choose_decoder(decoder):
-    assert decoder[:6] == 'deconv'
-    assert len(decoder)==7 
-
-    num_channels = 512
-    iheight, iwidth = 10, 8
-    kernel_size = int(decoder[6])
-    return DeConv(num_channels, kernel_size)
+class UpConv(Decoder):
+    # UpConv decoder consists of 4 upconv modules with decreasing number of channels and increasing feature map size
+    def upconv_module(self, in_channels):
+        # UpConv module: unpool -> 5*5 conv -> batchnorm -> ReLU
+        upconv = nn.Sequential(collections.OrderedDict([
+          ('unpool',    Unpool(in_channels)),
+          ('conv',      nn.Conv2d(in_channels,in_channels//2,kernel_size=5,stride=1,padding=2,bias=False)),
+          ('batchnorm', nn.BatchNorm2d(in_channels//2)),
+          ('relu',      nn.ReLU()),
+        ]))
+        return upconv
+
+    def __init__(self, in_channels):
+        super(UpConv, self).__init__()
+        self.layer1 = self.upconv_module(in_channels)
+        self.layer2 = self.upconv_module(in_channels//2)
+        self.layer3 = self.upconv_module(in_channels//4)
+        self.layer4 = self.upconv_module(in_channels//8)
+
+class UpProj(Decoder):
+    # UpProj decoder consists of 4 upproj modules with decreasing number of channels and increasing feature map size
+
+    class UpProjModule(nn.Module):
+        # UpProj module has two branches, with a Unpool at the start and a ReLu at the end
+        #   upper branch: 5*5 conv -> batchnorm -> ReLU -> 3*3 conv -> batchnorm 
+        #   bottom branch: 5*5 conv -> batchnorm
+
+        def __init__(self, in_channels):
+            super(UpProj.UpProjModule, self).__init__()
+            out_channels = in_channels//2
+            self.unpool = Unpool(in_channels)
+            self.upper_branch = nn.Sequential(collections.OrderedDict([
+              ('conv1',      nn.Conv2d(in_channels,out_channels,kernel_size=5,stride=1,padding=2,bias=False)),
+              ('batchnorm1', nn.BatchNorm2d(out_channels)),
+              ('relu',      nn.ReLU()),
+              ('conv2',      nn.Conv2d(out_channels,out_channels,kernel_size=3,stride=1,padding=1,bias=False)),
+              ('batchnorm2', nn.BatchNorm2d(out_channels)),
+            ]))
+            self.bottom_branch = nn.Sequential(collections.OrderedDict([
+              ('conv',      nn.Conv2d(in_channels,out_channels,kernel_size=5,stride=1,padding=2,bias=False)),
+              ('batchnorm', nn.BatchNorm2d(out_channels)),
+            ]))
+            self.relu = nn.ReLU()
+
+        def forward(self, x):
+            x = self.unpool(x)
+            x1 = self.upper_branch(x)
+            x2 = self.bottom_branch(x)
+            x = x1 + x2
+            x = self.relu(x)
+            return x
+
+    def __init__(self, in_channels):
+        super(UpProj, self).__init__()
+        self.layer1 = self.UpProjModule(in_channels)
+        self.layer2 = self.UpProjModule(in_channels//2)
+        self.layer3 = self.UpProjModule(in_channels//4)
+        self.layer4 = self.UpProjModule(in_channels//8)
+
+def choose_decoder(decoder, in_channels):
+    # iheight, iwidth = 10, 8
+    if decoder[:6] == 'deconv':
+        assert len(decoder)==7 
+        kernel_size = int(decoder[6])
+        return DeConv(in_channels, kernel_size)
+    elif decoder == "upproj":
+        return UpProj(in_channels)
+    elif decoder == "upconv":
+        return UpConv(in_channels)
+    else:
+        assert False, "invalid option for decoder: {}".format(decoder)
 
 
 class ResNet(nn.Module):
@@ -112,12 +190,12 @@ def __init__(self, layers, decoder, in_channels=3, out_channels=1, pretrained=Tr
         elif layers >= 50:
             num_channels = 2048
 
-        self.conv2 = nn.Conv2d(num_channels,512,kernel_size=1,bias=False)
-        self.bn2 = nn.BatchNorm2d(512)
-        self.decoder = choose_decoder(decoder)
+        self.conv2 = nn.Conv2d(num_channels,num_channels//2,kernel_size=1,bias=False)
+        self.bn2 = nn.BatchNorm2d(num_channels//2)
+        self.decoder = choose_decoder(decoder, num_channels//2)
 
         # setting bias=true doesn't improve accuracy
-        self.conv3 = nn.Conv2d(32,out_channels,kernel_size=3,stride=1,padding=1,bias=False)
+        self.conv3 = nn.Conv2d(num_channels//32,out_channels,kernel_size=3,stride=1,padding=1,bias=False)
         self.bilinear = nn.Upsample(size=(oheight, owidth), mode='bilinear')
 
         # weight init