问题描述
轨迹追踪模块,仅仅使用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)