code refactoring

Author: zisianw

data/config.py

@@ -2,11 +2,11 @@
 
 cfg = {
     'name': 'FaceBoxes',
-    'feature_maps': [[32, 32], [16, 16], [8, 8]],
-    'min_dim': 1024,
-    'steps': [32, 64, 128],
+    #'min_dim': 1024,
+    #'feature_maps': [[32, 32], [16, 16], [8, 8]],
+    # 'aspect_ratios': [[1], [1], [1]],
     'min_sizes': [[32, 64, 128], [256], [512]],
-    'aspect_ratios': [[1], [1], [1]],
+    'steps': [32, 64, 128],
     'variance': [0.1, 0.2],
     'clip': False,
     'loc_weight': 2.0,

layers/functions/prior_box.py

@@ -1,28 +1,21 @@
 import torch
 from itertools import product as product
 import numpy as np
+from math import ceil
 
 
 class PriorBox(object):
-    def __init__(self, cfg, box_dimension=None, image_size=None, phase='train'):
+    def __init__(self, cfg, image_size=None, phase='train'):
         super(PriorBox, self).__init__()
-        self.variance = cfg['variance']
+        #self.aspect_ratios = cfg['aspect_ratios']
         self.min_sizes = cfg['min_sizes']
         self.steps = cfg['steps']
-        self.aspect_ratios = cfg['aspect_ratios']
         self.clip = cfg['clip']
-        if phase == 'train':
-            self.image_size = (cfg['min_dim'], cfg['min_dim'])
-            self.feature_maps = cfg['feature_maps']
-        elif phase == 'test':
-            self.feature_maps = box_dimension.cpu().numpy().astype(np.int)
-            self.image_size = image_size
-        for v in self.variance:
-            if v <= 0:
-                raise ValueError('Variances must be greater than 0')
+        self.image_size = image_size
+        self.feature_maps = [[ceil(self.image_size[0]/step), ceil(self.image_size[1]/step)] for step in self.steps]
 
     def forward(self):
-        mean = []
+        anchors = []
         for k, f in enumerate(self.feature_maps):
             min_sizes = self.min_sizes[k]
             for i, j in product(range(f[0]), range(f[1])):
@@ -33,18 +26,18 @@ def forward(self):
                         dense_cx = [x*self.steps[k]/self.image_size[1] for x in [j+0, j+0.25, j+0.5, j+0.75]]
                         dense_cy = [y*self.steps[k]/self.image_size[0] for y in [i+0, i+0.25, i+0.5, i+0.75]]
                         for cy, cx in product(dense_cy, dense_cx):
-                            mean += [cx, cy, s_kx, s_ky]
+                            anchors += [cx, cy, s_kx, s_ky]
                     elif min_size == 64:
                         dense_cx = [x*self.steps[k]/self.image_size[1] for x in [j+0, j+0.5]]
                         dense_cy = [y*self.steps[k]/self.image_size[0] for y in [i+0, i+0.5]]
                         for cy, cx in product(dense_cy, dense_cx):
-                            mean += [cx, cy, s_kx, s_ky]
+                            anchors += [cx, cy, s_kx, s_ky]
                     else:
                         cx = (j + 0.5) * self.steps[k] / self.image_size[1]
                         cy = (i + 0.5) * self.steps[k] / self.image_size[0]
-                        mean += [cx, cy, s_kx, s_ky]
+                        anchors += [cx, cy, s_kx, s_ky]
         # back to torch land
-        output = torch.Tensor(mean).view(-1, 4)
+        output = torch.Tensor(anchors).view(-1, 4)
         if self.clip:
             output.clamp_(max=1, min=0)
         return output

layers/modules/multibox_loss.py

@@ -54,7 +54,7 @@ def forward(self, predictions, priors, targets):
                 shape: [batch_size,num_objs,5] (last idx is the label).
         """
 
-        loc_data, conf_data, _ = predictions
+        loc_data, conf_data = predictions
         priors = priors
         num = loc_data.size(0)
         num_priors = (priors.size(0))

models/faceboxes.py

@@ -27,20 +27,20 @@ def __init__(self):
     self.branch3x3_reduce_2 = BasicConv2d(128, 24, kernel_size=1, padding=0)
     self.branch3x3_2 = BasicConv2d(24, 32, kernel_size=3, padding=1)
     self.branch3x3_3 = BasicConv2d(32, 32, kernel_size=3, padding=1)
-  
+
   def forward(self, x):
     branch1x1 = self.branch1x1(x)
-    
+
     branch1x1_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1)
     branch1x1_2 = self.branch1x1_2(branch1x1_pool)
-    
+
     branch3x3_reduce = self.branch3x3_reduce(x)
     branch3x3 = self.branch3x3(branch3x3_reduce)
-    
+
     branch3x3_reduce_2 = self.branch3x3_reduce_2(x)
     branch3x3_2 = self.branch3x3_2(branch3x3_reduce_2)
     branch3x3_3 = self.branch3x3_3(branch3x3_2)
-    
+
     outputs = [branch1x1, branch1x1_2, branch3x3, branch3x3_3]
     return torch.cat(outputs, 1)
 
@@ -51,38 +51,38 @@ def __init__(self, in_channels, out_channels, **kwargs):
     super(CRelu, self).__init__()
     self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs)
     self.bn = nn.BatchNorm2d(out_channels, eps=1e-5)
-  
+
   def forward(self, x):
     x = self.conv(x)
     x = self.bn(x)
     x = torch.cat([x, -x], 1)
     x = F.relu(x, inplace=True)
     return x
-    
-    
+
+
 class FaceBoxes(nn.Module):
 
   def __init__(self, phase, size, num_classes):
     super(FaceBoxes, self).__init__()
     self.phase = phase
     self.num_classes = num_classes
     self.size = size
-    
+
     self.conv1 = CRelu(3, 24, kernel_size=7, stride=4, padding=3)
     self.conv2 = CRelu(48, 64, kernel_size=5, stride=2, padding=2)
-    
+
     self.inception1 = Inception()
     self.inception2 = Inception()
     self.inception3 = Inception()
-    
+
     self.conv3_1 = BasicConv2d(128, 128, kernel_size=1, stride=1, padding=0)
     self.conv3_2 = BasicConv2d(128, 256, kernel_size=3, stride=2, padding=1)
-    
+
     self.conv4_1 = BasicConv2d(256, 128, kernel_size=1, stride=1, padding=0)
     self.conv4_2 = BasicConv2d(128, 256, kernel_size=3, stride=2, padding=1)
-    
+
     self.loc, self.conf = self.multibox(self.num_classes)
-    
+
     if self.phase == 'test':
         self.softmax = nn.Softmax(dim=-1)
 
@@ -108,13 +108,12 @@ def multibox(self, num_classes):
     loc_layers += [nn.Conv2d(256, 1 * 4, kernel_size=3, padding=1)]
     conf_layers += [nn.Conv2d(256, 1 * num_classes, kernel_size=3, padding=1)]
     return nn.Sequential(*loc_layers), nn.Sequential(*conf_layers)
-    
+
   def forward(self, x):
-  
-    sources = list()
+
+    detection_sources = list()
     loc = list()
     conf = list()
-    detection_dimension = list()
 
     x = self.conv1(x)
     x = F.max_pool2d(x, kernel_size=3, stride=2, padding=1)
@@ -123,33 +122,28 @@ def forward(self, x):
     x = self.inception1(x)
     x = self.inception2(x)
     x = self.inception3(x)
-    detection_dimension.append(x.shape[2:])
-    sources.append(x)
+    detection_sources.append(x)
+
     x = self.conv3_1(x)
     x = self.conv3_2(x)
-    detection_dimension.append(x.shape[2:])
-    sources.append(x)
+    detection_sources.append(x)
+
     x = self.conv4_1(x)
     x = self.conv4_2(x)
-    detection_dimension.append(x.shape[2:])
-    sources.append(x)
-    
-    detection_dimension = torch.tensor(detection_dimension, device=x.device)
+    detection_sources.append(x)
 
-    for (x, l, c) in zip(sources, self.loc, self.conf):
+    for (x, l, c) in zip(detection_sources, self.loc, self.conf):
         loc.append(l(x).permute(0, 2, 3, 1).contiguous())
         conf.append(c(x).permute(0, 2, 3, 1).contiguous())
-        
+
     loc = torch.cat([o.view(o.size(0), -1) for o in loc], 1)
     conf = torch.cat([o.view(o.size(0), -1) for o in conf], 1)
 
     if self.phase == "test":
       output = (loc.view(loc.size(0), -1, 4),
-                self.softmax(conf.view(-1, self.num_classes)),
-                detection_dimension)
+                self.softmax(conf.view(-1, self.num_classes)))
     else:
       output = (loc.view(loc.size(0), -1, 4),
-                conf.view(conf.size(0), -1, self.num_classes),
-                detection_dimension)
-  
+                conf.view(conf.size(0), -1, self.num_classes))
+
     return output

test.py

@@ -113,13 +113,12 @@ def load_model(model, pretrained_path, load_to_cpu):
         scale = scale.to(device)
 
         _t['forward_pass'].tic()
-        out = net(img)  # forward pass
+        loc, conf = net(img)  # forward pass
         _t['forward_pass'].toc()
         _t['misc'].tic()
-        priorbox = PriorBox(cfg, out[2], (im_height, im_width), phase='test')
+        priorbox = PriorBox(cfg, image_size=(im_height, im_width))
         priors = priorbox.forward()
         priors = priors.to(device)
-        loc, conf, _ = out
         prior_data = priors.data
         boxes = decode(loc.data.squeeze(0), prior_data, cfg['variance'])
         boxes = boxes * scale / resize

train.py

@@ -9,6 +9,7 @@
 from layers.modules import MultiBoxLoss
 from layers.functions.prior_box import PriorBox
 import time
+import datetime
 import math
 from models.faceboxes import FaceBoxes
 
@@ -30,13 +31,19 @@
 if not os.path.exists(args.save_folder):
     os.mkdir(args.save_folder)
 
-img_dim = 1024
-rgb_means = (104, 117, 123) #bgr order
+img_dim = 1024 # only 1024 is supported
+rgb_mean = (104, 117, 123) # bgr order
 num_classes = 2
+num_gpu = args.ngpu
+num_workers = args.num_workers
 batch_size = args.batch_size
+momentum = args.momentum
 weight_decay = args.weight_decay
+initial_lr = args.lr
 gamma = args.gamma
-momentum = args.momentum
+max_epoch = args.max_epoch
+training_dataset = args.training_dataset
+save_folder = args.save_folder
 gpu_train = cfg['gpu_train']
 
 net = FaceBoxes('train', img_dim, num_classes)
@@ -58,17 +65,17 @@
         new_state_dict[name] = v
     net.load_state_dict(new_state_dict)
 
-if args.ngpu > 1 and gpu_train:
-    net = torch.nn.DataParallel(net, device_ids=list(range(args.ngpu)))
+if num_gpu > 1 and gpu_train:
+    net = torch.nn.DataParallel(net, device_ids=list(range(num_gpu)))
 
 device = torch.device('cuda:0' if gpu_train else 'cpu')
 cudnn.benchmark = True
 net = net.to(device)
 
-optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
+optimizer = optim.SGD(net.parameters(), lr=initial_lr, momentum=momentum, weight_decay=weight_decay)
 criterion = MultiBoxLoss(num_classes, 0.35, True, 0, True, 7, 0.35, False)
 
-priorbox = PriorBox(cfg)
+priorbox = PriorBox(cfg, image_size=(img_dim, img_dim))
 with torch.no_grad():
     priors = priorbox.forward()
     priors = priors.to(device)
@@ -79,10 +86,10 @@ def train():
     epoch = 0 + args.resume_epoch
     print('Loading Dataset...')
 
-    dataset = VOCDetection(args.training_dataset, preproc(img_dim, rgb_means), AnnotationTransform())
+    dataset = VOCDetection(training_dataset, preproc(img_dim, rgb_mean), AnnotationTransform())
 
-    epoch_size = math.ceil(len(dataset) / args.batch_size)
-    max_iter = args.max_epoch * epoch_size
+    epoch_size = math.ceil(len(dataset) / batch_size)
+    max_iter = max_epoch * epoch_size
 
     stepvalues = (200 * epoch_size, 250 * epoch_size)
     step_index = 0
@@ -95,15 +102,15 @@ def train():
     for iteration in range(start_iter, max_iter):
         if iteration % epoch_size == 0:
             # create batch iterator
-            batch_iterator = iter(data.DataLoader(dataset, batch_size, shuffle=True, num_workers=args.num_workers, collate_fn=detection_collate))
+            batch_iterator = iter(data.DataLoader(dataset, batch_size, shuffle=True, num_workers=num_workers, collate_fn=detection_collate))
             if (epoch % 10 == 0 and epoch > 0) or (epoch % 5 == 0 and epoch > 200):
-                torch.save(net.state_dict(), args.save_folder + 'FaceBoxes_epoch_' + repr(epoch) + '.pth')
+                torch.save(net.state_dict(), save_folder + 'FaceBoxes_epoch_' + str(epoch) + '.pth')
             epoch += 1
 
         load_t0 = time.time()
         if iteration in stepvalues:
             step_index += 1
-        lr = adjust_learning_rate(optimizer, args.gamma, epoch, step_index, iteration, epoch_size)
+        lr = adjust_learning_rate(optimizer, gamma, epoch, step_index, iteration, epoch_size)
 
         # load train data
         images, targets = next(batch_iterator)
@@ -112,33 +119,35 @@ def train():
 
         # forward
         out = net(images)
-        
+
         # backprop
         optimizer.zero_grad()
         loss_l, loss_c = criterion(out, priors, targets)
         loss = cfg['loc_weight'] * loss_l + loss_c
         loss.backward()
         optimizer.step()
         load_t1 = time.time()
-        print('Epoch:' + repr(epoch) + ' || epochiter: ' + repr(iteration % epoch_size) + '/' + repr(epoch_size) +
-              '|| Totel iter ' + repr(iteration) + ' || L: %.4f C: %.4f||' % (cfg['loc_weight']*loss_l.item(), loss_c.item()) +
-              'Batch time: %.4f sec. ||' % (load_t1 - load_t0) + 'LR: %.8f' % (lr))
+        batch_time = load_t1 - load_t0
+        eta = int(batch_time * (max_iter - iteration))
+        print('Epoch:{}/{} || Epochiter: {}/{} || Iter: {}/{} || L: {:.4f} C: {:.4f} || LR: {:.8f} || Batchtime: {:.4f} s || ETA: {}'.format(epoch, max_epoch, (iteration % epoch_size) + 1, epoch_size, iteration + 1, max_iter, loss_l.item(), loss_c.item(), lr, batch_time, str(datetime.timedelta(seconds=eta))))
 
-    torch.save(net.state_dict(), args.save_folder + 'Final_FaceBoxes.pth')
+
+    torch.save(net.state_dict(), save_folder + 'Final_FaceBoxes.pth')
 
 
 def adjust_learning_rate(optimizer, gamma, epoch, step_index, iteration, epoch_size):
-    """Sets the learning rate 
+    """Sets the learning rate
     # Adapted from PyTorch Imagenet example:
     # https://github.com/pytorch/examples/blob/master/imagenet/main.py
     """
-    if epoch < 0:
-        lr = 1e-6 + (args.lr-1e-6) * iteration / (epoch_size * 5) 
+    warmup_epoch = -1
+    if epoch <= warmup_epoch:
+        lr = 1e-6 + (initial_lr-1e-6) * iteration / (epoch_size * warmup_epoch)
     else:
-        lr = args.lr * (gamma ** (step_index))
+        lr = initial_lr * (gamma ** (step_index))
     for param_group in optimizer.param_groups:
         param_group['lr'] = lr
     return lr
-    
+
 if __name__ == '__main__':
     train()