TensorFlow 2.0 GPU Docker で CNN が動かない問題への対処 - 子育てしながらエンジニアしたい

以前の記事で、TensorFlow GPU Docker を Ubuntu 18.04 に導入しました。

edosha.hatenablog.jp

ここでは、とりあえず TensorFlow のチュートリアル Notebook が動作したので良しとしていたのですが、いざ実際の Notebook を動かそうとしたらエラーで動作しませんでした。
そのエラー内容と対処方法を記載します。

CNN を動かそうとすると cuDNN のエラーが出る

画像認識のタスクでは CNN を用いるのが一般的です。
以下のようなコードを使用していました。

from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.optimizers import SGD
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, GlobalAveragePooling2D, Dense

alpha = 0.5
input_shape = (224, 224, 3)
batchsize = 16

mobile = MobileNetV2(include_top=False, input_shape=input_shape, alpha=alpha, weights='imagenet')
x = mobile.output
x = GlobalAveragePooling2D()(x)
c = Dense(2, activation='softmax')(x)
model = Model(inputs=mobile.input,outputs=c)
model.compile(loss='categorical_crossentropy', optimizer=SGD(lr=5e-4, decay=0.00005), metrics=['accuracy'])
hist = model.fit(X_train, y_train, batch_size=batchsize, epochs=50, verbose=2)

別の GPU マシンでは問題なく動いていたのですが、TensorFlow 2.0 GPU Docker では Notebook に以下のようなエラーが出ていました。

Failed to get convolution algorithm. This is probably because cuDNN failed to initialize, so try looking to see if a warning log message was printed above.

コンソールの方には以下のようなエラーが出ていました。

 I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1006] successful NUMA node read from SysFS had negative value (-1),
 but there must be at least one NUMA node, so returning NUMA node zero
 I tensorflow/core/common_runtime/gpu/gpu_device.cc:1304] Created TensorFlow device (/job:localhost/replica:0/task:0/device
:GPU:0 with 7367 MB memory) -> physical GPU (device: 0, name: GeForce RTX 2070 SUPER, pci bus id: 0000:06:00.0, compute capability: 7.5)
 I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcudnn.so.7
 E tensorflow/stream_executor/cuda/cuda_dnn.cc:329] Could not create cudnn handle: CUDNN_STATUS_INTERNAL_ERROR

CUDNN_STATUS_INTERNAL_ERROR とは...???
TensorFlow のチュートリアルはちゃんと動かせていたので、これは CNN を動かすときだけの問題かもしれません。
いろいろ探し回って、ようやく以下の対処方法を見つけました。

原因と対処方法

見つけたのは Reddit の以下のスレッドです。

www.reddit.com

これによるとプリコンパイル済の TensorFlow に問題があるとか...
これって大問題のような気がするんですが対処してくれないんでしょうか...??

対処方法としては、以下のコードを model.fit() の前に追加すれば動作しました。

config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.compat.v1.Session(config=config)
tf.compat.v1.keras.backend.set_session(session)

この情報に辿り着くまでに 2 日ほどかかりました...