From 715aa0baaed2600517bb86c4ba0f922054898241 Mon Sep 17 00:00:00 2001 From: HonzaBenes Date: Thu, 12 Sep 2019 14:14:10 +0200 Subject: [PATCH] Using torch.no_grad() instead of deprecated parameter "volatile", specifying dim in softmax explicitly --- mtcnn/detector.py | 21 ++++++++++++--------- mtcnn/first_stage.py | 13 +++++++------ mtcnn/get_nets.py | 6 +++--- setup.py | 2 +- 4 files changed, 23 insertions(+), 19 deletions(-) diff --git a/mtcnn/detector.py b/mtcnn/detector.py index dde5f52..1c60c96 100644 --- 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 index 1ec857f..fd7d8b4 100644 --- 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 index 54a2c7a..fc86226 100644 --- 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 index f3f484a..7263327 100644 --- 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='',