[YOLO] 사용 장치 확인

YOLO에서 CPU를 사용 중인지 GPU를 사용 중인지 확인해 보자.

 

from ultralytics import YOLO

model = YOLO("yolo26n.pt")
print(f"Model is running on device: {next(model.model.parameters()).device}")

results = model.predict("bus.jpg", save=False, imgsz=320, conf=0.25)
#results = model.predict("bus.jpg", save=False, imgsz=320, conf=0.25, device='cpu')

print(f"Model is running on device: {next(model.model.parameters()).device}")

#results[0].show()

 

 

처음 모델을 로드하면 CPU를 사용한다. 하지만 예측 작업을 시작하면 자동으로 GPU(cuda)를 사용한다. (GPU 사용이 불가능한 경우 CPU를 사용한다)

 

from ultralytics import YOLO

model = YOLO("yolo26n.pt")
print(f"Model is running on device: {next(model.model.parameters()).device}")

#results = model.predict("bus.jpg", save=False, imgsz=320, conf=0.25)
results = model.predict("bus.jpg", save=False, imgsz=320, conf=0.25, device='cpu')

print(f"Model is running on device: {next(model.model.parameters()).device}")

#results[0].show()

 

 

예측 작업 시 device='cpu' 옵션을 추가하면 CPU를 사용한다.

 

※ 참고

Usage Example