admin/LC_NeoRefacer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Increased GPU usage

Browse Source
This commit is contained in:
Felipe Daragon
2025-05-15 22:30:23 +01:00
parent 59b6233882
commit badbcc6edf
9 changed files with 239 additions and 223 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
facelib/detection/yolov5face/utils/extract_ckpt.py
View File

@@ -1,5 +1,21 @@
import torch
import sys
sys.path.insert(0,'./facelib/detection/yolov5face')
model = torch.load('facelib/detection/yolov5face/yolov5n-face.pt', map_location='cpu')['model']
torch.save(model.state_dict(),'weights/facelib/yolov5n-face.pth')
import os
# Setup dynamic device selection
if torch.cuda.is_available():
device = torch.device("cuda")
elif torch.backends.mps.is_available():
device = torch.device("mps")
else:
device = torch.device("cpu")
sys.path.insert(0, './facelib/detection/yolov5face')
# Load the model to the selected device
ckpt = torch.load('facelib/detection/yolov5face/yolov5n-face.pt', map_location=device)
model = ckpt['model'].to(device)
# Save only the weights
os.makedirs('weights/facelib', exist_ok=True)
torch.save(model.state_dict(), 'weights/facelib/yolov5n-face.pth')