tf.config.experimental.get_memory_usage

Получение текущего использования памяти в байтах для выбранного устройства. (устаревшее)

tf.config.experimental.get_memory_usage(
    device
)

Эта функция устарела в пользу tf.config.experimental.get_memory_info. Вызов этой функции эквивалентен вызову tf.config.experimental.get_memory_info()['current'].

См. https://www.tensorflow.org/api_docs/python/tf/device для указания строк устройств.

Например:

gpu_devices = tf.config.list_physical_devices('GPU')
if gpu_devices:
  tf.config.experimental.get_memory_usage('GPU:0')

Не работает для ЦП.

Для графических процессоров TensorFlow по умолчанию будет выделять всю память, если не изменить это с помощью tf.config.experimental.set_memory_growth. Эта функция возвращает только память, которую фактически использует TensorFlow, а не память, которая была выделена на графическом процессоре TensorFlow.