Merge pull request #26 from Wuziyi616/master

[Enhance] Update code, pre-trained weight and future work of BiDet

README.md

@@ -1,14 +1,16 @@
 # BiDet
 This is the official pytorch implementation for paper: [*BiDet: An Efficient Binarized Object Detector*](https://arxiv.org/abs/2003.03961), which is accepted by CVPR2020. The code contains training and testing two binarized object detectors, SSD300 and Faster R-CNN, using our BiDet method on two datasets, PASCAL VOC and Microsoft COCO 2014.
 
+## Update
+- 2021.1: Our extended version of BiDet is accepted by T-PAMI! We further improve the performance of binary detectors and extend our method to multi model compression methods. Check it out [here](https://ieeexplore.ieee.org/abstract/document/9319565).
+- 2021.4.19: We provide BiDet-SSD300 pretrained weight on Pascal VOC dataset which achieves 66.0% mAP as described in the paper. You can download it [here](https://drive.google.com/file/d/1mURBX-EtoFanp-8wP6-n3u0E6Mq3enkv/view?usp=sharing).
+
 ## Quick Start
 ### Prerequisites
-- python 3.5+
+- python 3.6+
 - pytorch 1.0+
 - other packages include numpy, cv2, matplotlib, pillow, cython, cffi, msgpack, easydict, pyyaml
 
-**Note: **as [this issue](https://github.com/ZiweiWangTHU/BiDet/issues/17#issuecomment-725796486) pointed out, this repo is not compatible with PyTorch 1.7.0+. You can follow that instruction to modify the code and make it runnable using PyTorch 1.7.
-
 ### Dataset Preparation
 We conduct experiments on PASCAL VOC and Microsoft COCO 2014 datasets.  
 #### PASCAL VOC

ssd/eval_voc.py

@@ -428,7 +428,10 @@ def evaluate_detections(box_list, output_dir, dataset):
     if args.cuda:
         net = net.cuda()
         cudnn.benchmark = True
-    net.load_state_dict(torch.load(args.weight_path))
+    try:
+        net.load_state_dict(torch.load(args.weight_path))
+    except:
+        net.load_state_dict(torch.load(args.weight_path)['weight'])
     net.eval()
     print('Finished loading model!')
 
@@ -438,4 +441,5 @@ def evaluate_detections(box_list, output_dir, dataset):
                            VOCAnnotationTransform())
 
     # evaluation
-    test_net(net, dataset)
+    with torch.no_grad():
+        test_net(net, dataset)

ssd/layers/functions/detection.py

@@ -1,11 +1,9 @@
-import time
 import torch
-from torch.autograd import Function
 from ..box_utils import decode, nms
 from data import voc as cfg
 
 
-class Detect(Function):
+class Detect:
     """At test time, Detect is the final layer of   Decode location preds,
     apply non-maximum suppression to location predictions based on conf
     scores and threshold to a top_k number of output predictions for both
@@ -23,6 +21,9 @@ def __init__(self, num_classes, bkg_label, top_k=200, conf_thresh=0.03, nms_thre
         self.conf_thresh = conf_thresh
         self.variance = cfg['variance']
 
+    def __call__(self, loc_data, conf_data, prior_data):
+        return self.forward(loc_data, conf_data, prior_data)
+
     def forward(self, loc_data, conf_data, prior_data):
         """
         Args:
@@ -78,6 +79,9 @@ def __init__(self, num_classes, bkg_label, top_k=200, conf_thresh=0.03, nms_thre
         self.conf_thresh = conf_thresh
         self.variance = cfg['variance']
 
+    def __call__(self, loc_data, conf_data, prior_data, gt_class=None):
+        return self.forward(loc_data, conf_data, prior_data, gt_class)
+
     def forward(self, loc_data, conf_data, prior_data, gt_class=None):
         """
         Args:

ssd/train_bidet_ssd.py

@@ -74,7 +74,7 @@ def str2bool(v):
                     help='regularization loss weight for feature maps')
 parser.add_argument('--prior_weight', default=0., type=float,
                     help='loss weight for N(0, 1) prior')
-parser.add_argument('--sigma', default=0.001, type=float,
+parser.add_argument('--sigma', default=0., type=float,
                     help='scale factor controlling the sample procedure')
 parser.add_argument('--nms_conf_threshold', default=0.03, type=float,
                     help='confidence threshold for nms')
@@ -102,6 +102,7 @@ def str2bool(v):
 
 
 def train():
+    global REGULARIZATION_LOSS_WEIGHT, PRIOR_LOSS_WEIGHT, NMS_CONF_THRE
     if args.dataset == 'COCO':
         cfg = coco
         dataset = COCODetection(root=args.data_root,
@@ -118,9 +119,16 @@ def train():
     if args.cuda:
         cudnn.benchmark = True
 
+    opt_state_dict = None
     if args.resume:
         print('Resuming training, loading {}...'.format(args.weight_path))
-        ssd_net.load_weights(args.weight_path)
+        try:
+            ssd_net.load_state_dict(torch.load(args.weight_path))
+        except:  # checkpoint
+            print('Extracting from checkpoint')
+            ckp = torch.load(args.weight_path, map_location='cpu')
+            ssd_net.load_state_dict(ckp['weight'])
+            opt_state_dict = ckp['opt']
     else:
         if args.basenet.lower() != 'none':
             vgg_weights = torch.load(args.basenet)
@@ -138,6 +146,12 @@ def train():
                                weight_decay=args.weight_decay)
     else:
         exit(-1)
+    if opt_state_dict is not None:
+        print('Load optimizer state dict!')
+        optimizer.load_state_dict(opt_state_dict)
+        if get_lr(optimizer) != args.lr:
+            adjust_learning_rate(optimizer, args.lr)
+
     optimizer.zero_grad()
     criterion = MultiBoxLoss(cfg['num_classes'], 0.5, True, 0, True, 3, 0.5, False, args.cuda)
 
@@ -179,6 +193,7 @@ def train():
             epoch += 1
 
         if iteration in cfg['lr_steps']:
+            # add our BiDet loss in the after the first lr decay
             if step_index == 0:
                 args.reg_weight = 0.1
                 args.prior_weight = 0.2
@@ -345,7 +360,11 @@ def train():
             print('Saving state, iter:', iteration)
 
             loss_save = loc_loss_save + conf_loss_save + reg_loss_save + prior_loss_save
-            torch.save(net.module.state_dict(), logs_dir + '/model_' + str(iteration) +
+            checkpoint = {
+                'weight': net.module.state_dict(),
+                'opt': optimizer.state_dict()
+            }
+            torch.save(checkpoint, logs_dir + '/model_' + str(iteration) +
                        '_loc_' + str(round(loc_loss_save / 5000., 4)) +
                        '_conf_' + str(round(conf_loss_save / 5000., 4)) +
                        '_reg_' + str(round(reg_loss_save / 5000., 4)) +
@@ -371,15 +390,10 @@ def log_func(tensor):
     return tensor * torch.log(tensor)
 
 
-def adjust_learning_rate(optimizer, gamma, step):
-    """Sets the learning rate to the initial LR decayed by 10 at every
-        specified step
-    # Adapted from PyTorch Imagenet example:
-    # https://github.com/pytorch/examples/blob/master/imagenet/main.py
-    """
-    lr = args.lr * (gamma ** step)
+def adjust_learning_rate(optimizer, new_lr):
+    """Sets the learning rate of optimizer to new_lr."""
     for param_group in optimizer.param_groups:
-        param_group['lr'] = lr
+        param_group['lr'] = new_lr
 
 
 def get_lr(optimizer):