Skip to content

Commit

Permalink
optimize func: get_infer_gpuid (#13275)
Browse files Browse the repository at this point in the history
* optimize func: get_infer_gpuid

* apply code review suggestion
  • Loading branch information
GreatV authored Jul 9, 2024
1 parent c65a66c commit c95f6b2
Showing 1 changed file with 16 additions and 13 deletions.
29 changes: 16 additions & 13 deletions tools/infer/utility.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,19 +15,20 @@
import argparse
import os
import sys
import platform
import cv2
import numpy as np
import paddle
import PIL
from PIL import Image, ImageDraw, ImageFont
import math
from paddle import inference
import time
import random
from ppocr.utils.logging import get_logger


logger = get_logger()


def str2bool(v):
return v.lower() in ("true", "yes", "t", "y", "1")

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


def get_infer_gpuid():
sysstr = platform.system()
if sysstr == "Windows":
return 0
"""
Get the GPU ID to be used for inference.
Returns:
int: The GPU ID to be used for inference.
"""
if not paddle.device.is_compiled_with_rocm:
cmd = "env | grep CUDA_VISIBLE_DEVICES"
gpu_id_str = os.environ.get("CUDA_VISIBLE_DEVICES", "0")
else:
cmd = "env | grep HIP_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])
gpu_id_str = os.environ.get("HIP_VISIBLE_DEVICES", "0")

gpu_ids = gpu_id_str.split(",")
logger.warning(
"The first GPU is used for inference by default, GPU ID: {}".format(gpu_ids[0])
)
return int(gpu_ids[0])


def draw_e2e_res(dt_boxes, strs, img_path):
Expand Down

0 comments on commit c95f6b2

Please sign in to comment.