未验证 提交 09604c38 编写于 作者: M MissPenguin 提交者: GitHub

Merge pull request #4824 from ronny1996/cp_pr4823

cherry-pick #4823
...@@ -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():
cmd = "nvidia-smi" if not paddle.fluid.core.is_compiled_with_rocm():
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
cmd = "env | grep CUDA_VISIBLE_DEVICES" if not paddle.fluid.core.is_compiled_with_rocm():
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.
先完成此消息的编辑!
想要评论请 注册