こんにちは。過去の投稿をアップデートしたものです。2021年5月時点のDeepLearning環境構築方法を、NVIDIA GeForce RTX 3090 が搭載された Windows 10 に TensorFlow をインストールすることにより紹介します。
クラウドで試したい場合は、AzureのNCv3シリーズをご利用ください。

NVIDIA CUDA と cuDNN のインストール

まず、最新のNVIDIAドライバをインストールしてください。TensorFlow2.4の場合は、CUDA 11.2.2 と cuDNN 8.1.1 をNVIDIAのサイトからダウンロードしてください。

NVIDIAドライバダウンロード
https://www.nvidia.co.jp/Download/index.aspx?lang=jp
CUDA Toolkit 11.2.2 (March 2021)
https://developer.nvidia.com/cuda-toolkit-archive
The NVIDIA CUDA® Deep Neural Network library (cuDNN) for CUDA 11.2.2
https://developer.nvidia.com/rdp/cudnn-archive
（Downloadページから、「Download cuDNN v8.1.1 (Feburary 26th, 2021), for CUDA 11.0,11.1 and 11.2」を選択、次に「cuDNN Library for Windows 10」を選択してダウンロードを行います。）

初めにCUDAのインストーラーを実行してインストールします。CUDAをインストールした後にcuDNNをインストールします。ダウンロードしたcuDNNのzipを展開して、CUDAのインストールフォルダ（C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.2）にコピーします。CUDAのフォルダ名称に合わせてcuDNNのファイルをコピーしてください。なお、cuDNNのダウンロードには、ユーザの登録が必要になりますので、先に登録しておいてください。

※あと、TensorFlowをGPUで動かした時に、以下のようなメッセージで正しくDLLを読み込むことができない場合は、cusolver64_11.dllをcusolver64_10.dllとしてコピーすると動きます。

tensorflow/stream_executor/platform/default/dso_loader.cc:60] Could not load dynamic library 'cusolver64_10.dll'; dlerror: cusolver64_10.dll not found

コピー

copy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.2\bin\cusolver64_11.dll" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.2\bin\cusolver64_10.dll"

Pythonのインストール

プログラミング環境を整えていきます。機械学習をはじめとしたプログラミング言語としてはPythonが人気です。Python環境としてはAnacondaが有名で分析や機械学習を行うためのライブラリも問題なく利用できます。

Anacondaは下記のサイトよりダウンロードします。利用しているOSに合わせて適切なパッケージをインストールしてください。ボタンをそのままクリックし続ける（標準の設定）だけで大丈夫です。

Anaconda Download
https://www.continuum.io/downloads

Python の環境の作成

スタートメニューからAnaconda Promptを実行して、環境を作成していきます。環境を作ることで作業環境を分けることができます。

conda create -n py38 python=3.8

TensorFlow のインストール

TensorFlowをインストールします。TensorFlowは容易に複雑なネットワークを作ることができるDeep Learningのフレームワークです。

インストールする前に環境を先ほど作成した環境に変えておきます。

conda activate py38

次に、下記のコマンドでTensorFlowをインストールします。

pip install tensorflow==2.4

TensorFlow で GeForce RTX の機能を使ってみる

ソースコードは、以下のGitHubに mnist の Jupyter notebook として置いておきました。

https://github.com/KentaroAOKI/tf_tensorcore_test/blob/master/tensorcore_test-mixed_float16.ipynb
（https://www.tensorflow.org/guide/mixed_precision より)

混合精度

混合精度を使う場合、mixed_precision.set_policyで'mixed_float16'を設定します。以下の箇所で設定しています。
compute capability が 8.0 以上の GPU (RTX30xx系) では、float32 を使用してもパフォーマンスが変わらない可能性があります。これは、低精度の計算を自動的に使用するTensorFloat-32を使用しているからのようです。

policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_global_policy(policy)
print('Compute dtype: {}'.format(policy.compute_dtype))
print('Variable dtype: {}'.format(policy.variable_dtype))

Tensor Core

Tensor Coreは、次元を8の倍数にすることにより動作します。以下の箇所でDenseのユニット数を設定しています。

if tf.config.list_physical_devices('GPU'):
  print('The model will run with 4096 units on a GPU')
  num_units = 4096

Dense以外では、以下のように8の倍数を設定することによりTensorCoreが使われるようです。
https://www.tensorflow.org/guide/mixed_precision#gpu_performance_tips

tf.keras.layers.Dense(units=64)
tf.keras.layers.Conv2d(filters=48, kernel_size=7, stride=3)
tf.keras.layers.LSTM(units=64)
tf.keras.Model.fit(epochs=2, batch_size=128)

以下に、4096のDenseの結果を載せておきます。（RTX 3090）
最初のエポックは、モデルの最適化に時間を費やすため、遅くなる場合があります。

Epoch 1/5
6/6 [==============================] - 3s 137ms/step - loss: 4.9136 - accuracy: 0.2524 - val_loss: 0.8942 - val_accuracy: 0.8266
Epoch 2/5
6/6 [==============================] - 0s 33ms/step - loss: 0.7335 - accuracy: 0.8078 - val_loss: 0.3190 - val_accuracy: 0.9093
Epoch 3/5
6/6 [==============================] - 0s 30ms/step - loss: 0.4380 - accuracy: 0.8750 - val_loss: 0.2476 - val_accuracy: 0.9317
Epoch 4/5
6/6 [==============================] - 0s 29ms/step - loss: 0.2887 - accuracy: 0.9141 - val_loss: 0.3424 - val_accuracy: 0.9013
Epoch 5/5
6/6 [==============================] - 0s 29ms/step - loss: 0.2620 - accuracy: 0.9230 - val_loss: 0.2557 - val_accuracy: 0.9217
313/313 - 0s - loss: 0.2618 - accuracy: 0.9169
Test loss: 0.2617703676223755
Test accuracy: 0.9168999791145325

さいごに

GeForce RTX 3090 は速いですね！d(´∀`*)ではー。