WindowsにTensorFlow2.1 + GPU環境を構築する

cuDNNの初期化に失敗

Kerasを使っていてバックエンドにTensorFlowを使っているときにConv2Dレイヤーが含まれるモデルのトレーニングをしたところ以下のようなエラーが発生したので対応しました。
結果的にはCUDA ToolkitとcuDNNをインストールして対応しました。

This is probably because cuDNN failed to initialize windows

環境情報

発生時の環境です。

まさかゲーム以外のことにGeForce 2070を使うことになろうとは思いもしませんでした・・・。
GeForce RTX 2070を購入して浮かれているポストはROG-STRIX-RTX2070-O8G-GAMING買いましたになります。

CUDAのダウンロードとインストール

TensorFlowのバージョンによって使うCUDA Toolkitのバージョンを指定する必要があります。
今回はTensorFlow 2.1 を使うのでCUDA Toolkit 10.1 update2をダウンロードします。
CUDA Toolkit Download

ダウンロード後は適当にインストールします。
インストール後、ターミナルでnvcc -Vを実行して以下のような表示になればインストール完了です。

cuDNNのダウンロードとインストール

先程インストールしたCUDA Toolkitに対応したcuDNNをダウンロートしてインストールします。
今回はCuda Toolkit 10.1を使用するのでcuDNN v7.6.5 (November 5th, 2019), for CUDA 10.1をダウンロードします。

残念ながらインストーラーはないようなので、昔ながらの方法でPATHを通すかCUDA Toolkitのプログラムフォルダにファイルをコピーする必要があるみたいです。
今回はCUDA Toolkitのプログラムフォルダにコピーします。
ダウンロードしたzipファイルを解凍します。
すると{bin, lib, include}というフォルダがあるのでCUDA Toolkitのフォルダにコピーします。
自分の環境では以下のフォルダでした。
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.1
上書きされるものはなかったので安心です。

動作確認

エラーが発生した畳み込みのレイヤー（Conv2D）が含まれるモデルでトレーニングしてみてエラーが発生しなくなればOKです。

当然モデルの構築にもよりますが、自分の環境ではCPUだと1エポック約52秒かかっていた処理がGPUだと5秒にまで短縮されました。
処理時間が約1/10ですね。