Skip to content

Commit 1b60cf2

Browse files
authored
Merge pull request #1145 from common-io/master
FastMtCnn Use of Gpu for Face Detection
2 parents 357a397 + 131a7fb commit 1b60cf2

File tree

1 file changed

+4
-1
lines changed

1 file changed

+4
-1
lines changed

deepface/detectors/FastMtCnn.py

+4-1
Original file line numberDiff line numberDiff line change
@@ -62,13 +62,16 @@ def build_model(self) -> Any:
6262
# this is not a must dependency. do not import it in the global level.
6363
try:
6464
from facenet_pytorch import MTCNN as fast_mtcnn
65+
import torch
6566
except ModuleNotFoundError as e:
6667
raise ImportError(
6768
"FastMtcnn is an optional detector, ensure the library is installed."
6869
"Please install using 'pip install facenet-pytorch' "
6970
) from e
7071

71-
face_detector = fast_mtcnn(device="cpu")
72+
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
73+
face_detector = fast_mtcnn(device=device)
74+
7275
return face_detector
7376

7477

0 commit comments

Comments
 (0)