tensorflowでGPUメモリの使用量を抑えたいが上手くいかない

Question

tensorflowを実行するとGPUメモリの殆どを確保してしまうので、それを抑えたく、tfのバージョンに合わせて
tf.config.experimental.set_memory_growthや、gpu_options.allow_growth = True等を試しましたが、
どうしても大部分（3.7/4.0GB程度）が確保されてしまい、制限できません。

ちなみに、Jupyter Notebookを使用しています。

環境

• windows 10

• GPU : NVIDIA GeForce GTX 1650（専用メモリは4G）

• python 3.7.9 + TensorFlow-GPU 1.15.0 (python 3.6.12 + TensorFlow-GPU 2.0.0)

• Kerasはインストールしてない

• Visual Studio Community 2017

• CUDA 10.0

• cuDNN 7.6.0

---

tensorflowのメモリ使用量を抑制するコード（エラーは出ず、正常に実行されます）

tensorflow-gpu=2.0.0の環境と、tensorflow-gpu=1.15の環境を、Anacondaで作成し、
それぞれ実行しましたが、どうしてもメモリを抑えられません。

import tensorflow as tf
print("Tensorflow Ver: ", tf.__version__)

if tf.test.gpu_device_name():
    print('Default GPU Device: {}'.format(tf.test.gpu_device_name()))
else:
    print("Please install GPU version of TF")

gpu_id = 0
print(tf.__version__)
if tf.__version__ >= "2.1.0":
    physical_devices = tf.config.list_physical_devices('GPU')
    tf.config.list_physical_devices('GPU')
    tf.config.set_visible_devices(physical_devices[gpu_id], 'GPU')
    tf.config.experimental.set_memory_growth(physical_devices[gpu_id], True)

elif tf.__version__ >= "2.0.0":
    physical_devices = tf.config.experimental.list_physical_devices('GPU')
    tf.config.experimental.set_visible_devices(physical_devices[gpu_id], 'GPU')
    tf.config.experimental.set_memory_growth(physical_devices[gpu_id], True)

else:
    from tensorflow.keras.backend import set_session
    config = tf.ConfigProto(
        gpu_options=tf.GPUOptions(
            visible_device_list=str(gpu_id), # specify GPU number
            allow_growth=True
        )
    )
    set_session(tf.Session(config=config))

検証に使用したCNNモデル

from tensorflow.keras import datasets, layers, models

(train_images, train_labels), (test_images, test_labels) = datasets.cifar10.load_data()
train_images, test_images = train_images / 255.0, test_images / 255.0

model = models.Sequential()
model.add(layers.Conv2D(32, (3, 3), activation='relu', input_shape=(32, 32, 3))) #Conv2D(32, (3, 3)):２次元畳み込み層、「3ｘ3」の大きさのフィルタを32種類使用
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(64, (3, 3), activation='relu')) #Conv2D(64, (3, 3)):２次元畳み込み層、「3ｘ3」の大きさのフィルタを64種類使用
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(64, (3, 3), activation='relu'))

# Dense Layerの追加（全結合）- Classifier(分類)
model.add(layers.Flatten())
model.add(layers.Dense(units=64, activation='relu'))
model.add(layers.Dense(units=10)) #10個のアウトプット

model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])
model.fit(train_images, train_labels, batch_size=64, epochs=1, validation_data=(test_images, test_labels))

###　（参考）↓学習実行と同時にメモリを取りに行く

Answer 1

Anaconda Promptで
set TF_FORCE_GPU_ALLOW_GROWTH=true
python プログラムのファイル名.py
としたら、メモリー使用量が減りました

環境変数 TF_FORCE_GPU_ALLOW_GROWTH の設定をせずに、ご提示のコード(allow_growth=True とか実行している方)をプログラムのファイルの先頭に追加しても、効果はありませんでした

Comments

[koyamashinji]

毎度回答ありがとうございます。
（Jupyter Notebookを使っていますが）set TF_FORCE_GPU_ALLOW_GROWTH=true python プログラムのファイル名.ipynb
としても、やはりメモリを取りに行ってしまうようです・・・（tf2.0.0とtf1.15どちらの環境でも同じ結果となりました）
調べても原因がわからず・・Jupyterを使用していることに関係あるでしょうか。

[jbpb0]

ipynb ではなく、それを py に変換して、set ...=true を実行した直後に、同じアナコンダプロンプトで、
python ***.py
で実行してみてください
それなら、当方と同じ実行方法なので、おそらく効果あると思います

また、Windowsのデスクトップの「PCアイコン」を右クリック→プロパティ→システムの詳細設定→環境変数
で、「TF_FORCE_GPU_ALLOW_GROWTH」を新規作成して、値に「true」を設定してみてください
そうすれば、Jupyterでも効果あるかもしれません

[jbpb0]

ipynb ファイルに
import os
print(os.getenv('TF_FORCE_GPU_ALLOW_GROWTH'))
を追加したら、何て表示されますでしょうか?
true なら環境変数が正常に反映されてますが、None だと反映されてません

もし None なら、import tensorflow... より前に下記を追加してみるとか
import os
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'

[koyamashinji]

print(os.getenv('TF_FORCE_GPU_ALLOW_GROWTH'))の結果がNoneだった為、
import tensorflow... より前にimport os
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'を実行したところ、
tensorflowのバージョン1.15の環境では正常にメモリが制限されました！有難うございます。
（それでも1.2/4.0GB使ってしまうようですが・・・）
同じコードをtensorflow2.0.0の環境で実行した場合も、正常にメモリ制限されました！！

大変助かりました。毎度、有難うございます。感謝致します。

Answer 2

jbpb0様のご提案に基づき、以下のとおりコードを書くと、無事解決できました。

ポイントは、

• os.environ["TF_FORCE_GPU_ALLOW_GROWTH"]= "true"　を、import tensorflowの前に実行すること

正直、（★）以降のコードについては、かなり色んなサイトやブログで紹介されていますが
（★）以降のコードを実行せずとも、tensorflowをimportする前に　環境変数 TF_FORCE_GPU_ALLOW_GROWTHをTrueに設定すれば、正常にメモリ抑制できました。

# この部分は必ずTensorflowをimportする前に実行する
import os
os.environ["TF_FORCE_GPU_ALLOW_GROWTH"]= "true"　

# tensorflowをインポートし、バージョンを確認
import tensorflow as tf
print("Tensorflow Ver: ", tf.__version__)

# GPUを認識しているか確認
if tf.test.gpu_device_name():
    print('Default GPU Device: {}'.format(tf.test.gpu_device_name()))
else:
    print("Please install GPU version of TF")

# TFのメモリ制限が有効か確認（trueが返ったら有効、Noneが返ったら無効）
print(os.getenv('TF_FORCE_GPU_ALLOW_GROWTH'))

#（★）
# Tensorflowのバージョンに応じて、メモリ使用量を抑えるコードを実行
# gpu_id = 0
# print(tf.__version__)
# if tf.__version__ >= "2.1.0":
#    physical_devices = tf.config.list_physical_devices('GPU')
#    tf.config.list_physical_devices('GPU')
#    tf.config.set_visible_devices(physical_devices[gpu_id], 'GPU')
#    tf.config.experimental.set_memory_growth(physical_devices[gpu_id], True)

# elif tf.__version__ >= "2.0.0":
#    #TF2.0
#    physical_devices = tf.config.experimental.list_physical_devices('GPU')
#    tf.config.experimental.set_visible_devices(physical_devices[gpu_id], 'GPU')
#    tf.config.experimental.set_memory_growth(physical_devices[gpu_id], True)

# else:
#    from tensorflow.keras.backend import set_session
#    config = tf.ConfigProto(
#        gpu_options=tf.GPUOptions(
#            visible_device_list=str(gpu_id), # specify GPU number
#            allow_growth=True
#        )
#    )
#    set_session(tf.Session(config=config))

Comments

[jbpb0]

（★）以降のコードも要ります?

[koyamashinji]

（★）以降のコードを消して実行したところ、うまくいきましたので、要らないようです。ご指摘有難うございました。