Merge pull request #75 from n1mmy/master

Allow using different data types for MTCNN model.

Author: timesler

models/utils/detect_face.py

@@ -22,7 +22,8 @@ def detect_face(imgs, minsize, pnet, rnet, onet, threshold, factor, device):
 
     imgs = torch.as_tensor(imgs, device=device)
 
-    imgs = imgs.permute(0, 3, 1, 2).float()
+    model_dtype = next(pnet.parameters()).dtype
+    imgs = imgs.permute(0, 3, 1, 2).type(model_dtype)
 
     batch_size = len(imgs)
     h, w = imgs.shape[2:4]
@@ -178,7 +179,7 @@ def generateBoundingBox(reg, probs, scale, thresh):
     image_inds = mask_inds[:, 0]
     score = probs[mask]
     reg = reg[:, mask].permute(1, 0)
-    bb = mask_inds[:, 1:].float().flip(1)
+    bb = mask_inds[:, 1:].type(reg.dtype).flip(1)
     q1 = ((stride * bb + 1) / scale).floor()
     q2 = ((stride * bb + cellsize - 1 + 1) / scale).floor()
     boundingbox = torch.cat([q1, q2, score.unsqueeze(1), reg], dim=1)

tests/travis_test.py

@@ -146,6 +146,46 @@ def get_image(path, trans):
 mtcnn(img, save_path='data/tmp.png')
 
 
+#### MTCNN TYPES TEST ####
+
+img = Image.open('data/multiface.jpg')
+
+mtcnn = MTCNN(keep_all=True)
+boxes_ref, _ = mtcnn.detect(img)
+_ = mtcnn(img)
+
+mtcnn = MTCNN(keep_all=True).double()
+boxes_test, _ = mtcnn.detect(img)
+_ = mtcnn(img)
+
+box_diff = boxes_ref[np.argsort(boxes_ref[:,1])] - boxes_test[np.argsort(boxes_test[:,1])]
+total_error = np.sum(np.abs(box_diff))
+print('\nfp64 Total box error: {}'.format(total_error))
+
+assert total_error < 1e-2
+
+
+# half is not supported on CPUs, only GPUs
+if torch.cuda.is_available():
+
+    mtcnn = MTCNN(keep_all=True, device='cuda').half()
+    boxes_test, _ = mtcnn.detect(img)
+    _ = mtcnn(img)
+
+    box_diff = boxes_ref[np.argsort(boxes_ref[:,1])] - boxes_test[np.argsort(boxes_test[:,1])]
+    print('fp16 Total box error: {}'.format(np.sum(np.abs(box_diff))))
+
+    # test new automatic multi precision to compare
+    if hasattr(torch.cuda, 'amp'):
+        with torch.cuda.amp.autocast():
+            mtcnn = MTCNN(keep_all=True, device='cuda')
+            boxes_test, _ = mtcnn.detect(img)
+            _ = mtcnn(img)
+
+        box_diff = boxes_ref[np.argsort(boxes_ref[:,1])] - boxes_test[np.argsort(boxes_test[:,1])]
+        print('AMP total box error: {}'.format(np.sum(np.abs(box_diff))))
+
+            
 #### MULTI-IMAGE TEST ####
 
 mtcnn = MTCNN(keep_all=True)