Increased GPU usage
This commit is contained in:
@@ -11,7 +11,13 @@ from facelib.detection.retinaface.retinaface_net import FPN, SSH, MobileNetV1, m
|
||||
from facelib.detection.retinaface.retinaface_utils import (PriorBox, batched_decode, batched_decode_landm, decode, decode_landm,
|
||||
py_cpu_nms)
|
||||
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
|
||||
if torch.cuda.is_available():
|
||||
device = torch.device("cuda")
|
||||
elif torch.backends.mps.is_available():
|
||||
device = torch.device("mps")
|
||||
else:
|
||||
device = torch.device("cpu")
|
||||
|
||||
|
||||
def generate_config(network_name):
|
||||
@@ -367,4 +373,4 @@ class RetinaFace(nn.Module):
|
||||
# self.total_frame += len(frames)
|
||||
# print(self.batch_time / self.total_frame)
|
||||
|
||||
return final_bounding_boxes, final_landmarks
|
||||
return final_bounding_boxes, final_landmarks
|
||||
Reference in New Issue
Block a user