how to programmatically determine available GPU memory with tensorflow?

I actually found an answer in this old question of mine
. To bring some additional benefit to readers I tested the mentioned program

import nvidia_smi

nvidia_smi.nvmlInit()

handle = nvidia_smi.nvmlDeviceGetHandleByIndex(0)
# card id 0 hardcoded here, there is also a call to get all available card ids, so we could iterate

info = nvidia_smi.nvmlDeviceGetMemoryInfo(handle)

print("Total memory:", info.total)
print("Free memory:", info.free)
print("Used memory:", info.used)

nvidia_smi.nvmlShutdown()

on colab with the following result:

Total memory: 17071734784
Free memory: 17071734784
Used memory: 0

The actual GPU I had there was a Tesla P100 as can be seen from executing

!nvidia-smi

and observing the output

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 440.44       Driver Version: 418.67       CUDA Version: 10.1     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  Tesla P100-PCIE...  Off  | 00000000:00:04.0 Off |                    0 |
| N/A   32C    P0    26W / 250W |      0MiB / 16280MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

If you’re using tensorflow-gpu==2.5, you can use

tf.config.experimental.get_memory_info('GPU:0')

to get the actual consumed GPU memory by TF. Nvidia-smi tells you nothing, as TF allocates everything for itself and leaves nvidia-smi no information to track how much of that pre-allocated memory is actually being used.

In summary, the best solution that worked well is using: tf.config.experimental.get_memory_info('DEVICE_NAME')

This function returns a dictionary with two keys:

• ‘current’: The current memory used by the device, in bytes
• ‘peak’: The peak memory used by the device across the run of the program, in bytes.

The value of these keys is the ACTUAL memory used not the allocated one that is returned by nvidia-smi.

In reality, for GPUs, TensorFlow will allocate all the memory by default rendering using nvidia-smi to check for the used memory in your code useless.
Even if, tf.config.experimental.set_memory_growth is set to true, Tensorflow will no more allocate the whole available memory but is going to remain in allocating more memory than the one is used and in a discrete manner,
i.e. allocates 4589MiB then 8717MiB then 16943MiB then 30651 MiB, etc.

A small note concerning the get_memory_info() is that it doesn’t return correct values if used in a tf.function() decorated function. Thus, the peak key shall be used after executing tf.function() decorated function
to determine the peak memory used.

For older versions of Tensorflow, tf.config.experimental.get_memory_usage('DEVICE_NAME') was the only available function and only returned the used memory (no option for determining the peak memory).

Final note, you can also consider the Tensorflow Profiler available with Tensorboard to get information regarding your memory usage.

Hope this helps 🙂

sharing my over-engineered solution based on y.selivonchyk ‘s solution.

import os
import tempfile
import subprocess
import traceback
import pandas as pd

TH = 0.05

def get_one_available_gpu_device_id():
    gpu_device=-1
    try:
        with tempfile.TemporaryDirectory() as tmpdirname:
            fname = os.path.join(tmpdirname,'query.csv')
            cmd_list = f'nvidia-smi --format=csv --query-gpu=memory.total,memory.free,memory.used,pci.bus_id,index --filename={fname}'.split(' ')
            subprocess.check_output(cmd_list)
            if not os.path.exists(fname):
                raise ValueError("csv file not found")
            df = pd.read_csv(fname)
            df['gpu_mem_total']=df['memory.total [MiB]'].apply(lambda x: int(x.split(' ')[0]))
            df['gpu_mem_used']=df[' memory.used [MiB]'].apply(lambda x: int(x.split(' ')[1]))
            df['gpu_usage_prct']=df['gpu_mem_used']/df['gpu_mem_total']
            df['gpu_id']=df[' index']
            print(df)
            df = df.sort_values('gpu_usage_prct')
            avail = df[df.gpu_usage_prct < TH].reset_index()
            if len(avail)>0:
                gpu_device = avail.loc[0,'gpu_id']
    except:
        traceback.print_exc()

    return int(gpu_device)

gpu_device = get_one_available_gpu_device_id()
print(f'gpu_device {gpu_device}')

likely, you can then run a subprocess and specify the vacant gpu index:
f"CUDA_VISIBLE_DEVICES={gpu_device}"