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.
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. |
© 2020 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 3.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/versions/r2.4/api_docs/python/tf/config/experimental/get_memory_usage