torch.no_grad() instead of deprecated "volatile", dim set in softmax …

…call explicitly
TropComplique · Sep 12, 2019 · 0d5d23b · 0d5d23b
1 parent b79915d
commit 0d5d23b
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 19 deletions.
diff --git a/mtcnn/detector.py b/mtcnn/detector.py
@@ -76,10 +76,11 @@ def detect_faces(image, min_face_size=20.0,
     # STAGE 2
 
     img_boxes = get_image_boxes(bounding_boxes, image, size=24)
-    img_boxes = Variable(torch.FloatTensor(img_boxes), volatile=True)
-    output = rnet(img_boxes)
-    offsets = output[0].data.numpy()  # shape [n_boxes, 4]
-    probs = output[1].data.numpy()  # shape [n_boxes, 2]
+    with torch.no_grad():
+        img_boxes = Variable(torch.FloatTensor(img_boxes))
+        output = rnet(img_boxes)
+        offsets = output[0].data.numpy()  # shape [n_boxes, 4]
+        probs = output[1].data.numpy()  # shape [n_boxes, 2]
 
     keep = np.where(probs[:, 1] > thresholds[1])[0]
     bounding_boxes = bounding_boxes[keep]
@@ -97,11 +98,13 @@ def detect_faces(image, min_face_size=20.0,
     img_boxes = get_image_boxes(bounding_boxes, image, size=48)
     if len(img_boxes) == 0: 
         return [], []
-    img_boxes = Variable(torch.FloatTensor(img_boxes), volatile=True)
-    output = onet(img_boxes)
-    landmarks = output[0].data.numpy()  # shape [n_boxes, 10]
-    offsets = output[1].data.numpy()  # shape [n_boxes, 4]
-    probs = output[2].data.numpy()  # shape [n_boxes, 2]
+
+    with torch.no_grad():
+        img_boxes = Variable(torch.FloatTensor(img_boxes))
+        output = onet(img_boxes)
+        landmarks = output[0].data.numpy()  # shape [n_boxes, 10]
+        offsets = output[1].data.numpy()  # shape [n_boxes, 4]
+        probs = output[2].data.numpy()  # shape [n_boxes, 2]
 
     keep = np.where(probs[:, 1] > thresholds[2])[0]
     bounding_boxes = bounding_boxes[keep]

diff --git a/mtcnn/first_stage.py b/mtcnn/first_stage.py
@@ -31,12 +31,13 @@ def run_first_stage(image, net, scale, threshold):
     img = cv2.resize(image, (sw, sh))
     img = np.asarray(img, 'float32')
 
-    img = Variable(torch.FloatTensor(_preprocess(img)), volatile=True)
-    output = net(img)
-    probs = output[1].data.numpy()[0, 1, :, :]
-    offsets = output[0].data.numpy()
-    # probs: probability of a face at each sliding window
-    # offsets: transformations to true bounding boxes
+    with torch.no_grad():
+        img = Variable(torch.FloatTensor(_preprocess(img)))
+        output = net(img)
+        probs = output[1].data.numpy()[0, 1, :, :]
+        offsets = output[0].data.numpy()
+        # probs: probability of a face at each sliding window
+        # offsets: transformations to true bounding boxes
 
     boxes = _generate_bboxes(probs, offsets, scale, threshold)
     if len(boxes) == 0:

diff --git a/mtcnn/get_nets.py b/mtcnn/get_nets.py
@@ -68,7 +68,7 @@ def forward(self, x):
         x = self.features(x)
         a = self.conv4_1(x)
         b = self.conv4_2(x)
-        a = F.softmax(a)
+        a = F.softmax(a, dim=1)
         return b, a
 
 
@@ -113,7 +113,7 @@ def forward(self, x):
         x = self.features(x)
         a = self.conv5_1(x)
         b = self.conv5_2(x)
-        a = F.softmax(a)
+        a = F.softmax(a, dim=1)
         return b, a
 
 
@@ -166,5 +166,5 @@ def forward(self, x):
         a = self.conv6_1(x)
         b = self.conv6_2(x)
         c = self.conv6_3(x)
-        a = F.softmax(a)
+        a = F.softmax(a, dim=1)
         return c, b, a
diff --git a/setup.py b/setup.py
@@ -2,7 +2,7 @@
 
 setup(
     name='mtcnn-pytorch',
-    version='1.0.0',
+    version='1.0.2',
     packages=['mtcnn'],
     package_data={'mtcnn': ['weights/*.npy']},
     url='',