support pytorch 1.0.1

Author: zisianw

data/data_augment.py

@@ -47,7 +47,7 @@ def _crop(image, boxes, labels, img_dim):
         boxes_t[:, 2:] = np.minimum(boxes_t[:, 2:], roi[2:])
         boxes_t[:, 2:] -= roi[:2]
 
-        # keep the cropped image after resizing has at less one face > 16 pixel
+	# make sure that the cropped image contains at least one face > 16 pixel at training image scale
         b_w_t = (boxes_t[:, 2] - boxes_t[:, 0] + 1) / w * img_dim
         b_h_t = (boxes_t[:, 3] - boxes_t[:, 1] + 1) / h * img_dim
         mask_b = np.minimum(b_w_t, b_h_t) > 16.0
@@ -186,7 +186,6 @@ def __init__(self, img_dim, rgb_means):
         self.rgb_means = rgb_means
 
     def __call__(self, image, targets):
-        image = image.astype(np.float32)
         assert targets.shape[0] > 0, "this image does not have gt"
 
         boxes = targets[:, :-1].copy()

layers/modules/multibox_loss.py

@@ -70,9 +70,6 @@ def forward(self, predictions, priors, targets):
         if GPU:
             loc_t = loc_t.cuda()
             conf_t = conf_t.cuda()
-        # wrap targets
-        loc_t = Variable(loc_t, requires_grad=False)
-        conf_t = Variable(conf_t, requires_grad=False)
 
         pos = conf_t > 0
 
@@ -81,7 +78,7 @@ def forward(self, predictions, priors, targets):
         pos_idx = pos.unsqueeze(pos.dim()).expand_as(loc_data)
         loc_p = loc_data[pos_idx].view(-1, 4)
         loc_t = loc_t[pos_idx].view(-1, 4)
-        loss_l = F.smooth_l1_loss(loc_p, loc_t, size_average=False)
+        loss_l = F.smooth_l1_loss(loc_p, loc_t, reduction='sum')
 
         # Compute max conf across batch for hard negative mining
         batch_conf = conf_data.view(-1, self.num_classes)
@@ -101,7 +98,7 @@ def forward(self, predictions, priors, targets):
         neg_idx = neg.unsqueeze(2).expand_as(conf_data)
         conf_p = conf_data[(pos_idx+neg_idx).gt(0)].view(-1,self.num_classes)
         targets_weighted = conf_t[(pos+neg).gt(0)]
-        loss_c = F.cross_entropy(conf_p, targets_weighted, size_average=False)
+        loss_c = F.cross_entropy(conf_p, targets_weighted, reduction='sum')
 
         # Sum of losses: L(x,c,l,g) = (Lconf(x, c) + αLloc(x,l,g)) / N
         N = max(num_pos.data.sum().float(), 1)

train.py

@@ -4,7 +4,6 @@
 import torch.optim as optim
 import torch.backends.cudnn as cudnn
 import argparse
-from torch.autograd import Variable
 import torch.utils.data as data
 from data import AnnotationTransform, VOCDetection, detection_collate, preproc, cfg
 from layers.modules import MultiBoxLoss
@@ -108,12 +107,8 @@ def train():
 
         # load train data
         images, targets = next(batch_iterator)
-        if gpu_train:
-            images = Variable(images.cuda())
-            targets = [Variable(anno.cuda()) for anno in targets]
-        else:
-            images = Variable(images)
-            targets = [Variable(anno) for anno in targets]
+        images = images.to(device)
+        targets = [anno.to(device) for anno in targets]
 
         # forward
         out = net(images)