From 29e4b248b1ac29ddd8f7b6de9e4b3d23b0c00c40 Mon Sep 17 00:00:00 2001 From: Nick Butts Date: Fri, 1 Mar 2024 05:47:54 -0600 Subject: [PATCH] Fixed the metric so it works with ROCM librariers. --- yolox/utils/metric.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/yolox/utils/metric.py b/yolox/utils/metric.py index 506b58281..6618ab998 100644 --- a/yolox/utils/metric.py +++ b/yolox/utils/metric.py @@ -26,12 +26,21 @@ def get_total_and_free_memory_in_Mb(cuda_device): "nvidia-smi --query-gpu=memory.total,memory.used --format=csv,nounits,noheader" ) devices_info = devices_info_str.read().strip().split("\n") - if "CUDA_VISIBLE_DEVICES" in os.environ: - visible_devices = os.environ["CUDA_VISIBLE_DEVICES"].split(',') - cuda_device = int(visible_devices[cuda_device]) - total, used = devices_info[int(cuda_device)].split(",") - return int(total), int(used) + if len(devices_info) > 1: + if "CUDA_VISIBLE_DEVICES" in os.environ: + visible_devices = os.environ["CUDA_VISIBLE_DEVICES"].split(',') + cuda_device = int(visible_devices[cuda_device]) + total, used = devices_info[int(cuda_device)].split(",") + else: + devices_info_str = os.popen( + "rocm-smi --showmeminfo vram --csv" + ) + devices_info = devices_info_str.read().strip().split("\n") + _, total, used = devices_info[1].split(',') + total = int(total) // (1024 * 1024) + used = int(used) // (1024 * 1024) + return int(total), int(used) def occupy_mem(cuda_device, mem_ratio=0.9): """