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')

Не работает для CPU.

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