Deep Learningとメモリー・マネジメント

 今回は我々の永遠の課題であるメモリー管理の話です。

 Kaggleのようなコンペティションでもそうですが、我々がDeep Learningのコードを実行する際には利用可能なメモリー容量の制限の範囲内でコードを実行する必要があります。

1. メモリー上に格納せずディスクに書き込む

 これは処理速度が遅くなるのであまりオススメしませんが、どうしても使用できるメモリー容量が少ない場合は変数としてメモリーに格納しておかず、デスクに書き込む。なお、書き込んだ後で下記に触れる「gc.collect()」で変数用に使ったメモリー領域を開放してあげる必要があります。

2. 再利用しない途中生成データはこまめに消す

 コードの中で変数を用意してデータをそこに格納し、処理を行った後にそれを再利用しない場合はその変数自体を削除し、gc.collect()をコールして、削除した変数で使ったメモリー領域を開放します。なお、gc.collect()にはメモリー開放のタイミングを具体的に設定するオプションもありますが、経験上、単純にgc.collect()することで十分と思います。

import gc

# delete used parameter
del parameter_name

# release memory space for above
gc.collect()

3. TensorFlowがGPUメモリー（VRAM）を少しずつで確保するように

 最近のNvidiaの民生用グラフィックカード（RTX20x0シリーズ）に実装されているVRAMは8GBです。個人的には納得いきません。このご時世に8GBではCVPRを行ったりGANを使う上ではちょっときつい。RTX2080シリーズでは16GB、RTX2070シリーズでは12GB、そしてRTX2060シリーズでは8GBというようにプロセッサのグレードとVRAM容量を同時に増減させて欲しいです。

 という個人的な話はさておき、TensorFlowでGPUに処理をさせる際、デフォルトではありったけのVRAMを確保しようとします。これでメモリー・リークが発生して処理ができなくなるのですが、この「最初にありったけ確保」をやめさせて、必要なだけ徐々に確保する「Memory Glowth」を行わせるには「tf.config.experimental.set_memory_growth(gpu, True)」をコールします。

import tensorflow as tf

gpus = tf.config.experimental.list_physical_devices('GPU')

if gpus:
    
    try:
        # Currently, memory growth needs to be the same across GPUs
        for gpu in gpus:

            # GPUの占有メモリーを一気に確保せず、処理に必要な量を随時に確保する
            tf.config.experimental.set_memory_growth(gpu, True)

            # コードを実行するプロセッサを指定したにもかかわらず、それが存在しない場合は自動で指定を変える
            tf.config.set_soft_device_placement(True)
            tf.debugging.set_log_device_placement(True)

        logical_gpus = tf.config.experimental.list_logical_devices('GPU')
        print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
        
    except RuntimeError as e:
        
        # Memory growth must be set before GPUs have been initialized
        print(e)

4. 予約占有されたGPUメモリー（VRAM）を開放する

 TensorFlowの処理が完了したにも関わらず、GPUのVRAMが相変わらずパンパンなままになっていることってありますよね。そのような場合に強制的にTensorFlowが確保したVRAM領域を開放するには下記のようにコールします。

from numba import cuda

# 下記の場合は/gpu:0に関してCUDA処理をcloseしています。
cuda.select_device(0)
cuda.close()

 2020年の中盤には次期RTXシリーズの「RTX30x0」がリリースされるようです。噂ではVRAMが大きくなり、値段はRTX20x0シリーズよりもお安くなるようです。RTX3080TIまで上位にならない機種でVRAMが12-16GBになることを祈って今回は終わります。