tf.config.experimental.get_memory_usage

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

tf.config.experimental.get_memory_usage(
    device
)

См. 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')

Не работает для процессора.