yolov8轨迹追踪算法无法指定gpu

原创已于 2024-08-03 14:26:26 修改 · 431 阅读

1 ·

CC 4.0 BY-SA版权

文章标签：

#YOLO #python #机器学习 #人工智能 #算法 #pytorch

于 2024-08-03 10:35:50 首次发布

问题描述

轨迹追踪模块，仅仅使用demo推理，可以通过 device参数指定加载模型所使用的GPU，但有时会失效。

class yolov8Tracker(baseTracker):
    def __init__(self):
        super(yolov8Tracker, self).__init__()
        self.init_model()

    def init_model(self):
        self.weights = DETECTOR_PATH
        self.device = gpu if torch.cuda.is_available() else 'cpu'
        self.model = YOLO(self.weights)
        self.model.to(self.device)
        self.m = self.model
        self.names = self.model.module.names if hasattr(self.model, 'module') else self.model.names

    def track(self, im, polys=None):
        res = self.model.track(im, tracker="bytetrack.yaml", persist=True, imgsz=self.img_size, conf=self.conf,
                               iou=self.iou, device=1, verbose=False)

我在开发项目时，发现这种调用方式，有时无法顺利指定期望使用的GPU。

解决方法

通过torch.device的方法传参

之后测试问题解决。

from ultralytics import YOLO
from test.utils import is_point_in_roi
import logging
import torch
try:
    import os

    gpu = int(os.environ["gpu"])
except Exception as e:
    gpu = 0
DETECTOR_PATH = 'weights/yolov8n.pt'
logging.info(f"track_detect.py gpu: {gpu}")

logging.info(os.environ['CUDA_VISIBLE_DEVICES'])

def tracker_infer(weight_path):
    model = YOLO(weight_path)

    def process(frame):
        res = model.track(frame, tracker="bytetrack.yaml", device=torch.device(f"cuda:{gpu}"), verbose=False)