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 выделил на графическом процессоре.