View source on GitHub |
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. |