From 613c571ee8930580e5a6ed466747b206150cdf60 Mon Sep 17 00:00:00 2001 From: joizhang Date: Wed, 22 Dec 2021 14:52:10 +0800 Subject: [PATCH 1/2] feat: facenet pytorch support cpu. --- app/recognition/facenet.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/app/recognition/facenet.py b/app/recognition/facenet.py index e6a698d..09ead31 100644 --- a/app/recognition/facenet.py +++ b/app/recognition/facenet.py @@ -32,8 +32,8 @@ class FaceNetTorch(FaceRecognition): def create_model(self): # set TORCH_HOME in your os environment self.model = InceptionResnetV1(pretrained='vggface2') - assert torch.cuda.is_available() - self.model.cuda() + if torch.cuda.is_available(): + self.model.cuda() self.model.eval() def get_embeddings(self, input_image: numpy.ndarray): -- Gitee From 6d809b82a64f5dcdda6fdd54ad8db92c3c84cdf9 Mon Sep 17 00:00:00 2001 From: joizhang Date: Wed, 22 Dec 2021 14:53:29 +0800 Subject: [PATCH 2/2] fix: facenet pytorch support cpu. --- app/recognition/facenet.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/app/recognition/facenet.py b/app/recognition/facenet.py index 09ead31..349a1fc 100644 --- a/app/recognition/facenet.py +++ b/app/recognition/facenet.py @@ -38,7 +38,9 @@ class FaceNetTorch(FaceRecognition): def get_embeddings(self, input_image: numpy.ndarray): x = self.transform(input_image) - x = x.unsqueeze(0).cuda() + x = x.unsqueeze(0) + if torch.cuda.is_available(): + x = x.cuda() embeddings = self.model(x) embeddings = embeddings.detach().cpu().numpy() return embeddings[0] -- Gitee