TensorFlow, PyTorchで共通のCUDA, Cudnnを利用する

・PyTorch
PyTorchのCuda対応バージョンを確認。
2022/3現在、10.2と11.3。以下では11.3を例に説明。
https://pytorch.org/


・TensorFlow
「テスト済みのビルド構成」で上記CUDAバージョンより新しくて、最も近いバージョンを選択。
CUDA11.3の例であれば、tensorflow-2.6.0のCUDA11.2が該当。
https://www.tensorflow.org/install/source#common_installation_problems

 

・CUDA
PyTorchに対応しているCUDAをインストール(11.3)
https://developer.nvidia.com/cuda-toolkit-archive


・Cudnn

インストールしたCUDAに対応するCudnnをインストール(8.2.1)
https://developer.nvidia.com/rdp/cudnn-archive