Как оптимизировать код PyTorch и Ultralytics Yolo для использования графического процессора? ⇐ Python

[Anonymous]

Я работаю над проектом, который предполагает обнаружение и отслеживание объектов. Для обнаружения объектов я использую yolov8, а для отслеживания — трекер SORT. После запуска приведенного ниже кода загрузка моего графического процессора всегда ниже 10%, а загрузка ЦП всегда превышает 40%. Я установил cuda, cudnn и установил torch с помощью cuda. Я также скомпилировал opencv с поддержкой cuda. Я использую RTX 4060 ti, но похоже, что он не используется.
Есть ли способ дальнейшей оптимизации приведенного ниже кода, чтобы вся работа выполнялась графическим процессором и не процессор?
from src.sort import *
import cv2
import time
import torch
import numpy as np
from ultralytics import YOLO

device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"Using device: {device}")
sort_tracker = Sort(max_age=20, min_hits=2, iou_threshold=0.05)
model = YOLO('yolov8s.pt').to(device)

cap = cv2.VideoCapture(0)

while True:
ret, frame = cap.read()
if not ret:
print("**No frame received**")
continue

results = model(frame)
dets_to_sort = np.empty((0, 6))
for result in results:
for obj in result.boxes:
bbox = obj.xyxy[0].cpu().numpy().astype(int)
x1, y1, x2, y2 = bbox

conf = obj.conf.item()
class_id = int(obj.cls.item())
dets_to_sort = np.vstack((dets_to_sort, np.array([x1, y1, x2, y2, conf, class_id])))

tracked_dets = sort_tracker.update(dets_to_sort)
for det in tracked_dets:
x1, y1, x2, y2 = [int(i) for i in det[:4]]
track_id = int(det[8]) if det[8] is not None else 0
class_id = int(det[4])
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 4)
cv2.putText(frame, f"{track_id}", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 255, 255), 3)

frame = cv2.resize(frame, (800, int(frame.shape[0] * 800 / frame.shape[1])), interpolation=cv2.INTER_NEAREST)
cv2.imshow("Frame", frame)
key = cv2.waitKey(1)
if key == ord("q"):
break
if key == ord("p"):
cv2.waitKey(-1)

cap.release()
cv2.destroyAllWindows()


Подробнее здесь: https://stackoverflow.com/questions/786 ... tilize-gpu