From 05bdd0abfeea6230af822a58b10c858d6d93bf4b Mon Sep 17 00:00:00 2001 From: luanhailiang Date: Fri, 22 Sep 2023 20:50:12 +0800 Subject: [PATCH] device_count --- app.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/app.py b/app.py index 36f823e..c38a584 100644 --- a/app.py +++ b/app.py @@ -1,6 +1,7 @@ import cv2 import json import time +import torch import requests import threading import numpy as np @@ -61,6 +62,8 @@ def worker(id,source,region=None,stream=False): region = 10 else: region = int(region) + device_count = torch.cuda.device_count() + print("GPU device count:", device_count) print("start loading model...",id,source,region) model = YOLO('./best.pt') print("start loaded model!!!") @@ -92,7 +95,7 @@ def worker(id,source,region=None,stream=False): print("Error opening video stream....") break continue - result = model.track(frame,show=False,stream=False,persist=True,device=int(id)%4) + result = model.track(frame,show=False,stream=False,persist=True,device=int(id)%device_count) use += " track:"+str((time.time()*1000) - bgn) del(ret) del(frame)