提交 96a3f90c 编写于 作者: L LDOUBLEV

add gpu check

上级 74e0d114
......@@ -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
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册