From 96a3f90ced9b7c39bccea04cc1b2fd35293871c1 Mon Sep 17 00:00:00 2001
From: LDOUBLEV <liuvv0203@outlook.com>
Date: Tue, 17 Aug 2021 02:29:41 +0000
Subject: [PATCH] add gpu check

---
 tools/infer/predict_det.py | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/tools/infer/predict_det.py b/tools/infer/predict_det.py
index 3de00d83..6a9ae987 100755
--- a/tools/infer/predict_det.py
+++ b/tools/infer/predict_det.py
@@ -101,6 +101,7 @@ class TextDetector(object):
         if args.benchmark:
             import auto_log
             pid = os.getpid()
+            gpu_id = self.get_infer_gpuid()
             self.autolog = auto_log.AutoLogger(
                 model_name="det",
                 model_precision=args.precision,
@@ -110,13 +111,26 @@ class TextDetector(object):
                 inference_config=self.config,
                 pids=pid,
                 process_name=None,
-                gpu_ids=0,
+                gpu_ids=gpu_id,
                 time_keys=[
                     'preprocess_time', 'inference_time', 'postprocess_time'
                 ],
                 warmup=2,
                 logger=logger)
 
+    def get_infer_gpuid(self):
+        cmd = "nvidia-smi"
+        res = os.popen(cmd).readlines()
+        if len(res) == 0:
+            return None
+        cmd = "env | grep CUDA_VISIBLE_DEVICES"
+        env_cuda = os.popen(cmd).readlines()
+        if len(env_cuda) == 0:
+            return 0
+        else:
+            gpu_id = env_cuda[0].strip().split("=")[1]
+            return int(gpu_id[0])
+
     def order_points_clockwise(self, pts):
         """
         reference from: https://github.com/jrosebr1/imutils/blob/master/imutils/perspective.py
-- 
GitLab