Get the memory usage, in bytes, for the chosen device.
tf.config.experimental.get_memory_usage(
device
)
See https://www.tensorflow.org/api_docs/python/tf/device for specifying device
strings.
For example:
gpu_devices = tf.config.list_physical_devices('GPU')
if gpu_devices:
tf.config.experimental.get_memory_usage('GPU:0')
Does not work for CPU.
Args |
device
|
Device string to get the bytes in use for.
|
Returns |
Total memory usage in bytes.
|
Raises |
ValueError
|
Non-existent or CPU device specified.
|