提交 0a9c9f10 编写于 作者: R ronny1996

cherry-pick #4823

上级 ccf749e3
...@@ -20,6 +20,7 @@ import numpy as np ...@@ -20,6 +20,7 @@ import numpy as np
import json import json
from PIL import Image, ImageDraw, ImageFont from PIL import Image, ImageDraw, ImageFont
import math import math
import paddle
from paddle import inference from paddle import inference
import time import time
from ppocr.utils.logging import get_logger from ppocr.utils.logging import get_logger
...@@ -286,11 +287,17 @@ def create_predictor(args, mode, logger): ...@@ -286,11 +287,17 @@ def create_predictor(args, mode, logger):
def get_infer_gpuid(): def get_infer_gpuid():
if not paddle.fluid.core.is_compiled_with_rocm():
cmd = "nvidia-smi" cmd = "nvidia-smi"
else:
cmd = "rocm-smi"
res = os.popen(cmd).readlines() res = os.popen(cmd).readlines()
if len(res) == 0: if len(res) == 0:
return None return None
if not paddle.fluid.core.is_compiled_with_rocm():
cmd = "env | grep CUDA_VISIBLE_DEVICES" cmd = "env | grep CUDA_VISIBLE_DEVICES"
else:
cmd = "env | grep HIP_VISIBLE_DEVICES"
env_cuda = os.popen(cmd).readlines() env_cuda = os.popen(cmd).readlines()
if len(env_cuda) == 0: if len(env_cuda) == 0:
return 0 return 0
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册