TensorFlow (v2.2.0)でGPUを使う方法 (Windows)

概要

TensorFlowでGPUを使えるように設定したときのメモ。最新バージョン(2020年5月29日現在v2.2.0)での情報が公式以外あまり見当たらず少し苦戦してしまったので、動いた一例として載せておきます。
<参考ページ>
TensorFlow GPUサポート
CUDA Installation Guide for Microsoft Windows

環境

OS: Windows10 (64bit)
GPU: GeForce GTX 1050
Python: 3.8.2 (Anaconda3環境 ¹)

インストールするものリストとそのバージョン

Visual Studio 2019 Community
CUDA Toolkit v10.1 update2
[CuDNN v7.6.5 for CUDA 10.1] (https://developer.nvidia.com/rdp/cudnn-download) (ダウンロードにはNVIDIAの無料の会員登録が必要)
Tensorflow v2.2.0 (ただしpipを使いました¹)

やったこととその順番

1.Visual Studio のインストール

公式ページからインストーラーをダウンロードして実行するだけです。

2. CUDA Toolkit のインストール

TensorFlow公式ページで指定されているバージョン(v10.1)をインストール。v10.1の中で最新のupdate2を選びました。なおインストーラーのデフォルトでは、NVIDIA GPUドライバなど各種必要なものもCUDA Toolkitをインストールするのと同時にインストールされます。
ちなみに、異なるバージョンのCUDA Toolkitが既にインストールされている場合はエラーが出る可能性があります。その場合はNVIDIA関係のソフトをすべてアンインストールしてからインストールを行うとスムーズかもしれません。²

3. CUDA Toolkit が正常に入ったか確認

CUDAの公式インストールガイドには色々書いてあるのですが、インストール後にサンプルコードが走るか確認せよとのことだったので、一応確認しました。今回は各種あるサンプルの中で、テストとしてお勧めされていたbandwidthTestを使います。
C:\ProgramData\NVIDIA Corporation\CUDA Samples\v10.1\1_Utilities\bandwidthTest
にあるbandwidthTest_vs201X.slnをVisual Studioで開き(XはインストールされているVisual Studioに合わせてください。)、ビルドを実行します。すると、以下のディレクトリにbandwidthTest.exeが出力されます。

C:\ProgramData\NVIDIA Corporation\CUDA Samples\v10.1\bin\win64\Debug

あとはコマンドプロンプトを開いてexeファイルを実行し、以下のような出力になれば成功です。

C:\ProgramData\NVIDIA Corporation\CUDA Samples\v10.1\bin\win64\Debug>bandwidthTest.exe
[CUDA Bandwidth Test] - Starting...
Running on...

 Device 0: GeForce GTX 1050
 Quick Mode

 Host to Device Bandwidth, 1 Device(s)
 PINNED Memory Transfers
   Transfer Size (Bytes)        Bandwidth(GB/s)
   32000000                     6.1

 Device to Host Bandwidth, 1 Device(s)
 PINNED Memory Transfers
   Transfer Size (Bytes)        Bandwidth(GB/s)
   32000000                     6.2

 Device to Device Bandwidth, 1 Device(s)
 PINNED Memory Transfers
   Transfer Size (Bytes)        Bandwidth(GB/s)
   32000000                     95.7

Result = PASS

NOTE: The CUDA Samples are not meant for performance measurements. Results may vary when GPU Boost is enabled.

3. CuDNN のダウンロードおよび解凍

CuDNNのダウンロードには、NVIDIA Developer Program Membershipなるものが必要ですが、無料なので登録します。ログイン後、ダウンロードページ　(https://developer.nvidia.com/rdp/cudnn-download) から、CUDAのバージョンに合うものをダウンロードします(今回の場合、CuDNN v7.6.5 for CUDA 10.1)。
ダウンロードしたzipファイルを解凍すると、cudaという名前のフォルダが出てくるので、これを適当なフォルダに置きます。今回は公式にならって C:\tools\cuda とします。

4. PATHを通す

コマンドプロンプトを開いて以下を実行します。

SET PATH=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.1\bin;%PATH%
SET PATH=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.1\extras\CUPTI\libx64;%PATH%
SET PATH=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.1\include;%PATH%
SET PATH=C:\tools\cuda\bin;%PATH%

5. Tensorflow (v2.2.0) のインストール ¹

Anacondaでpython3の新しい必要最低限のものしか入っていない環境を作り、そこで
pip install tensorflow
を実行すればその時点での最新の安定バージョン(v2.2.0かそれ以上)がインストールされます。
(tensorflow v1系統ではここがtensorflow-gpuと分かれているのですがv2から統合されたみたいです。)

TensorFlowがGPU認識してるか確認

いよいよ確認のときです。

import tensorflow as tf

tf.test.is_built_with_cuda()
print(len(tf.config.experimental.list_physical_devices('GPU')))

を実行し、TrueとGPUの数(1以上)が返ってくれば成功です。

あとは普通にtensorflowを使えば早くなってるはず！

1. Anacondaでpipを使うのは背徳感があります。anaconda ディストリビューションにもtensorflow v2.1.0があったので、最初はcondaインストールしようとしたのですが、ここに来て、CUDA Toolkit v10.0でないとだめですとanacondaに怒られてしまい、やむを得ずpipにしました。今後condaインストールしようという方は、anacondaのtensorflowが対応しているCUDAのバージョンをチェックしてください。（但し、CUDAさえダウングレードすればcondaインストールでちゃんと動いたかなどは確認していませんのでご了承ください。) ↩ ↩² ↩³

2. 私は最初不注意でv10.2をインストールしてしまい、Toolkit だけアンインストールしてv10.1をインストールしようとした際にエラーが出ました。 ↩