请教一下
用 gpu 运行下面代码的时候，发现永远运行在 device 0 上
ocr = PaddleOCR(use_gpu=True, ...)
ocr.ocr(img)
是不是因为这段代码
PaddleOCR/tools/infer/utility.py
Lines 161 to 167 in 8b3b787
| | ifargs.use_gpu: |
| | gpu_id=get_infer_gpuid() |
| | ifgpu_idisNone: |
| | raiseValueError( |
| | "Not found GPU in current device. Please check your device or set args.use_gpu as False" |
| | ) |
| | config.enable_use_gpu(args.gpu_mem, 0) |
设置 device_id 永远是 0 导致的？
可以有什么方法手动设置 device 吗？谢谢
2条答案
按热度按时间nkhmeac61#
可以有什么方法手动设置 device 吗？谢谢
需要设置环境变量CUDA_VISIBLE_DEVICES，比如在1号GPU上运行
export CUDA_VISIBLE_DEVICES=1
python inference.py ....
vwkv1x7d2#
设置 CUDA_VISIBLE_DEVICES=1 确实可以让 get_infer_gpuid() 返回 1
但是如上面贴的代码 (PaddleOCR/tools/infer/utility.py:167)
config.enable_use_gpu(args.gpu_mem, 0)
这里参数写死了使用的 device 是 0，这会导致始终使用 0 卡