21 lines
596 B
Python
21 lines
596 B
Python
import torch
|
|
import sys
|
|
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') |