GPUのメモリ節約法(tensorflow)

こんにちは、にわかです。

テーマは、タイトル通りなのですが、GPUのメモリ(VRAM)使い果たし対策です。

初期化時のGPUのメモリ確保を制限する。

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  try:
    for gpu in gpus:
        tf.config.experimental.set_memory_growth(gpu, True) 
  except RuntimeError as e:
    print(e)

バッチサイズを1つにしたのに、画像の縦横の大きさを小さくしたのに、メモリの使い果たしをした人は上の方法で解決できるはずです。

※プログラムの開始時に、GPUのメモリを確保できるだけ確保するため(何かしらのプロセス用に)、後から学習データをGPUにロードする際にメモリの使い果たしで、ランタイムエラーを起こす場合があります。上のコードを書けば、必要になった分だけGPUのメモリ確保を行います。

ちなみに、CPUには上のコード使えないので注意してください。