実行確認環境
- Windows 10, 2004
- python 3.6.10
- tensorflow 2.3
- CUDA 10.1 update 2
- cuDNN 7.6.5
メモリの必要分確保
TF-1ではGPUのメモリはあるだけ確保する仕様がデフォルトでした。
そのため、GPUを必要な分だけ確保するおまじないを書くのが通例でした。
そのおまじないもTF-2.3で以下のように変わりました。
gpu_number
番目のGPUについて、必要な分だけのメモリを確保します。
import tensorflow as tf
def allocate_gpu_memory(gpu_number=0):
physical_devices = tf.config.experimental.list_physical_devices('GPU')
if physical_devices:
try:
print("Found {} GPU(s)".format(len(physical_devices)))
tf.config.set_visible_devices(physical_devices[gpu_number], 'GPU')
tf.config.experimental.set_memory_growth(physical_devices[gpu_number], True)
print("#{} GPU memory is allocated".format(gpu_number))
except RuntimeError as e:
print(e)
else:
print("Not enough GPU hardware devices available")
また、tensorflow2.0 + kerasでGPUメモリの使用量を抑える方法 によると、
tensorflow2.3からはset_memory_growthが不要になったようです。ドキュメント読んでないので詳しくは分からないですがモデル作って学習してもメモリを全部もっていったりはされなくなっていました。
とのことですが、残念ながら私の環境では上記内容は再現できませんでした。
確かに、デフォルトで複数枚のGPUメモリ全てを確保することはなくなりましたが、1枚のGPUメモリは全て確保されてしまいました。
さらに、全部のGPUメモリを数百MBくらい、薄く確保していました。
詳しい方、どなたか教えてください。
デバッグログ非表示
TF-2.xではGPU関係のログが大量に出てきます。
以下のようなログが最初に数十行くらいパーッと出てきて正直鬱陶しいです。
2020-08-21 21:57:32.240995: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library cudart64_101.dll
2020-08-21 21:57:34.374408: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library nvcuda.dll
2020-08-21 21:57:34.453727: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1716] Found device 0 with properties:
....
以下のようにすれば TF-2.3でログが全く表示されなくなりました。
TF-1.xでのログ非表示方法は使えなくなってますのでご注意を。
import os
import warnings
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
warnings.simplefilter(action='ignore', category=FutureWarning)
warnings.simplefilter(action='ignore', category=Warning)
import tensorflow as tf
tf.get_logger().setLevel('INFO')
tf.autograph.set_verbosity(0)
import logging
tf.get_logger().setLevel(logging.ERROR)
(21/7/2追記)
silence_tensorflow
というTF2ログを消すためだけのライブラリがあるようです。
こちらの方が簡単ですね。
TF2.3, 2.4で動作確認しました。
https://pypi.org/project/silence-tensorflow/
$ pip install silence-tensorflow
from silence_tensorflow import silence_tensorflow
silence_tensorflow()
import tensorflow as tf
TF-2.0 ~ TF-2.2 は使わない方がよい?
きちんと調査したわけではありませんが、TF-2.0 ~ TF-2.2はメモリリークの噂があるようです。
よほどの理由がない限りは TF-2.3 を使った方が無難でしょう。
cf. Memory leak - Issue #33309