github/deepface
1
1
Fork 0
mirror of https://github.com/serengil/deepface.git synced 2025-06-07 03:55:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: cuda change

Browse Source
This commit is contained in:
Shivam Singhal 2024-08-28 08:10:05 +05:30
parent e9d980fdd4
commit fa0ce75a5f
1 changed files with 7 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
deepface/models/facial_recognition/Facenet.py
View File

@ -1,7 +1,7 @@
import os import os
from typing import List, Any
import gdown import gdown
import numpy as np import numpy as np
from typing import List, Any
from deepface.commons import package_utils, folder_utils from deepface.commons import package_utils, folder_utils
from deepface.models.FacialRecognition import FacialRecognition from deepface.models.FacialRecognition import FacialRecognition
from deepface.commons.logger import Logger from deepface.commons.logger import Logger
@ -1742,15 +1742,17 @@ def load_facenet512d_onnx_model(
model (Any) model (Any)
""" """
try: try:
import torch # https://stackoverflow.com/questions/75267445
import onnxruntime as ort import onnxruntime as ort
except ModuleNotFoundError as e: except ModuleNotFoundError as e:
raise ImportError( raise ImportError(
"FaceNet512dONNX is an optional model, ensure the library is installed. " "FaceNet512ONNX is an optional model, ensure the library is installed. "
"Please install using 'pip install onnxruntime' or 'pip install onnxruntime-gpu' to use gpu" "Please install using 'pip install onnxruntime' or "
"'pip install onnxruntime-gpu' to use gpu"
) from e ) from e
if ort.get_device() == "GPU": if torch.cuda.is_available():
logger.info(f"using onnx GPU for inference") logger.info("using onnx GPU for inference")
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
else: else:
providers = ['CPUExecutionProvider'] providers = ['CPUExecutionProvider']