RTX3090/3080でCUDA11とcuDNN8に更新したらTensorflowが早くなった

RTX 3090を手に入れたが、どうにも遅い。

旧世代のTITAN RTXの方が1epochあたりの処理時間が早く、なんでかなーと思っていた。

結論から言うと、NVIDIA DRIVERは更新したけどCUDAやcuDNNを更新していなかったというだけ。

まずはnvidia-smiを見てみる。

右上にCUDA 11.1 と書いてあるので、ドライバーを更新したらCUDAも勝手に更新されたもんだと勘違いしていた。

遅いとはいえ、なまじTensorflowが動いていただけに気づくのが遅れてしまった。

実際のCUDAのバージョンを確認する。

$ nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2019 NVIDIA Corporation
Built on Sun_Jul_28_19:07:16_PDT_2019
Cuda compilation tools, release 10.1, V10.1.243

ドライバーは更新できたけど、CUDAそのものは10.1のままだった。

11.1に更新する。

$ wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/cuda-ubuntu1804.pin
$ sudo mv cuda-ubuntu1804.pin /etc/apt/preferences.d/cuda-repository-pin-600
$ sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/7fa2af80.pub
$ sudo add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/ /"
$ sudo apt-get update
$ sudo apt-get -y install cuda

ubuntu1804の所は自分のOSに合わせて書き換える。

$ nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2020 NVIDIA Corporation
Built on Mon_Oct_12_20:09:46_PDT_2020
Cuda compilation tools, release 11.1, V11.1.105
Build cuda_11.1.TC455_06.29190527_0

無事にCUDAが更新された。

次に、Ampereに対応したTensorflowを2.4に更新する。

Adds Support for TensorFloat-32 on Ampere based GPUs. TensorFloat-32, or TF32 for short, is a math mode for NVIDIA Ampere based GPUs and is enabled by default.
https://github.com/tensorflow/tensorflow/releases/tag/v2.4.0

pip install -U tensorflow

Tensorflowを更新してGPUが認識しなくなってしまった場合はcuDNNも更新する。

TensorFlow pip packages are now built with CUDA11 and cuDNN 8.0.2.
https://github.com/tensorflow/tensorflow/releases/tag/v2.4.0

$ wget https://developer.download.nvidia.com/compute/cuda/repos/${OS}/x86_64/cuda-${OS}.pin 

$ sudo mv cuda-${OS}.pin /etc/apt/preferences.d/cuda-repository-pin-600
$ sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/${OS}/x86_64/7fa2af80.pub
$ sudo add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/${OS}/x86_64/ /"
$ sudo apt-get update

Where ${OS} is ubuntu1804 or ubuntu2004.

これでTITAN RTXよりもRTX3090の方が早くなり、GPUの買い替えを実感することができた。

参考

CUDA Toolkit 12.1 Downloads

Get the latest feature updates to NVIDIA's proprietary compute stack.

Release TensorFlow 2.4.0 · tensorflow/tensorflow

Release 2.4.0 Major Features and Improvements tf.distribute introduces experimental support for asynchronous training of...

Overview — NVIDIA cuDNN 9.0.0 documentation