Skip to content

Commit c95f6b2

Browse files
authored
optimize func: get_infer_gpuid (#13275)
* optimize func: get_infer_gpuid * apply code review suggestion
1 parent c65a66c commit c95f6b2

File tree

1 file changed

+16
-13
lines changed

1 file changed

+16
-13
lines changed

tools/infer/utility.py

Lines changed: 16 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -15,19 +15,20 @@
1515
import argparse
1616
import os
1717
import sys
18-
import platform
1918
import cv2
2019
import numpy as np
2120
import paddle
2221
import PIL
2322
from PIL import Image, ImageDraw, ImageFont
2423
import math
2524
from paddle import inference
26-
import time
2725
import random
2826
from ppocr.utils.logging import get_logger
2927

3028

29+
logger = get_logger()
30+
31+
3132
def str2bool(v):
3233
return v.lower() in ("true", "yes", "t", "y", "1")
3334

@@ -333,20 +334,22 @@ def get_output_tensors(args, mode, predictor):
333334

334335

335336
def get_infer_gpuid():
336-
sysstr = platform.system()
337-
if sysstr == "Windows":
338-
return 0
337+
"""
338+
Get the GPU ID to be used for inference.
339339
340+
Returns:
341+
int: The GPU ID to be used for inference.
342+
"""
340343
if not paddle.device.is_compiled_with_rocm:
341-
cmd = "env | grep CUDA_VISIBLE_DEVICES"
344+
gpu_id_str = os.environ.get("CUDA_VISIBLE_DEVICES", "0")
342345
else:
343-
cmd = "env | grep HIP_VISIBLE_DEVICES"
344-
env_cuda = os.popen(cmd).readlines()
345-
if len(env_cuda) == 0:
346-
return 0
347-
else:
348-
gpu_id = env_cuda[0].strip().split("=")[1]
349-
return int(gpu_id[0])
346+
gpu_id_str = os.environ.get("HIP_VISIBLE_DEVICES", "0")
347+
348+
gpu_ids = gpu_id_str.split(",")
349+
logger.warning(
350+
"The first GPU is used for inference by default, GPU ID: {}".format(gpu_ids[0])
351+
)
352+
return int(gpu_ids[0])
350353

351354

352355
def draw_e2e_res(dt_boxes, strs, img_path):

0 commit comments

Comments
 (0)