TensorFlow(2系)で実際に使用しているGPUのメモリ量を確認する

はじめに

TensorFlowはGPUを利用する際、特にtf.config.experimental.set_memory_growthを設定しなければ、メモリ効率化のために全てのメモリを割り当てようとしてしまう。
そのためnvidida-smi等ではその割り当てられた割り当てられた量しか確認できず、内部では実際にどれだけのGPUメモリを使用しているのか確認する方法はないか知りたかった。

実行環境

Ubuntu: 20.04
TensorFlow: 2.4.0

方法

結論としてはtf.config.experimental.get_memory_usageを使えばよい。
GPUのみにしか使えない。

gpu_devices = tf.config.list_physical_devices('GPU')
if gpu_devices:
  tf.config.experimental.get_memory_usage('GPU:0')

上記のように、GPU：IDみたいに指定すると、実際に使用したbyteが返ってくる。

もし、間違っている点やより良い方法があれば教えてください。

参考

https://www.tensorflow.org/versions/r2.4/api_docs/python/tf/config/experimental/get_memory_usage
今回はversion2.4で試したが、TensorFlow2.6ではget_memory_usageの方が推奨らしい
https://www.tensorflow.org/api_docs/python/tf/config/experimental/get_memory_info
https://www.tensorflow.org/api_docs/python/tf/config/experimental/get_memory_usage