conda install のみでTensorFlowとGPU環境（CUDA、cuDNN）を構築する

TensorFlowでGPUを使って学習するために

TensorFlowはバージョン毎にCUDAやcuDNNを整備する必要があります。（複数のCUDA・cuDNNを一つのマシンに共存させることにより、TensorFlowが自動で環境を選択することもできるようですが、バグの温床になりそうなのでやっていません。。。）
この為、TensorFlowのバージョンとCUDA・cuDNNの対応を調べて、それらをインストールすることをずっとやっていました。。。そして以前、Anacondaの仮想環境にCUDAとcuDNNをインストールして、それを利用するという記事を書いたところ、@hatsuyuki396様から以下のコメントを頂きました。

conda install tensorflow=2.8.2=gpu_py39hc0c9373_0
conda install tensorflow=2.8.*=gpu_*
のようにビルドを指定すればcudatoolkitとcudnnが付属したtensorflowを入手できますよ
ビルドはconda searchで探せます

このコメントを頂いた時の心境としては「ガチ！？」という感じでした。

実際に試してみた

ということで、実際にコメント通りの方法で環境を構築できるのかを試してみました。一応、新たに仮想環境を作成して、そこにインストールするという形にしています。

conda create -n tf282 tensorflow=2.8.2=gpu_py39hc0c9373_0

各種インストールが終わり、GPUをTensorFlow上で確認したところ

GPUを確認するコード

from tensorflow.python.client import device_lib
device_lib.list_local_devices()

出力

[name: "/device:CPU:0"
device_type: "CPU"
memory_limit: 268435456
locality {
}
incarnation: 12142248063336448492
xla_global_id: -1
, name: "/device:GPU:0"
device_type: "GPU"
memory_limit: 10849157120
locality {
  bus_id: 1
  links {
  }
}
incarnation: 17153162280940221000
physical_device_desc: "device: 0, name: NVIDIA GeForce GTX 1080 Ti, pci bus id: 0000:01:00.0, compute capability: 6.1"
xla_global_id: 416903419
]

なんと、しっかりと認識していました！
念の為、実際に簡単なmnistの学習を回してみましたが、しっかりとGPUを利用して学習していました！今までの労力は一体。。。（PyTorchと似たような環境構築方法ですね）

conda search

一応、conda search からビルドを入手する手段を記載しておきます。

conda search tensorflow

と、打てば

出力

（省略）
tensorflow                     2.6.0 eigen_py37h34b007a_0  pkgs/main           
tensorflow                     2.6.0 eigen_py38hcc1cb13_0  pkgs/main           
tensorflow                     2.6.0 eigen_py39h4b72145_0  pkgs/main           
tensorflow                     2.6.0 mkl_py37h9d15365_0  pkgs/main           
tensorflow                     2.6.0 mkl_py38h874c8e8_0  pkgs/main           
tensorflow                     2.6.0 mkl_py39haac40d1_0  pkgs/main           
tensorflow                     2.8.2 eigen_py310h2d010dc_0  pkgs/main           
tensorflow                     2.8.2 eigen_py37h02982d4_0  pkgs/main           
tensorflow                     2.8.2 eigen_py38h4fd5193_0  pkgs/main           
tensorflow                     2.8.2 eigen_py39h2e3267a_0  pkgs/main           
tensorflow                     2.8.2 gpu_py310hf783e45_0  pkgs/main           
tensorflow                     2.8.2 gpu_py37hbb399c4_0  pkgs/main           
tensorflow                     2.8.2 gpu_py38h75b8afa_0  pkgs/main           
tensorflow                     2.8.2 gpu_py39hc0c9373_0  pkgs/main           
tensorflow                     2.8.2 mkl_py310hd2b8f8c_0  pkgs/main           
tensorflow                     2.8.2 mkl_py37h98100e6_0  pkgs/main           
tensorflow                     2.8.2 mkl_py38hb41d75a_0  pkgs/main           
tensorflow                     2.8.2 mkl_py39ha986a27_0  pkgs/main           
tensorflow                     2.9.1 eigen_py310h7bda0dd_0  pkgs/main           
tensorflow                     2.9.1 eigen_py37h054527f_0  pkgs/main           
tensorflow                     2.9.1 eigen_py38h6be4119_0  pkgs/main           
tensorflow                     2.9.1 eigen_py39h0984c0c_0  pkgs/main           
tensorflow                     2.9.1 gpu_py310h710b358_0  pkgs/main           
tensorflow                     2.9.1 gpu_py37hecb99ea_0  pkgs/main           
tensorflow                     2.9.1 gpu_py38h731a0ae_0  pkgs/main           
tensorflow                     2.9.1 gpu_py39hd4ae112_0  pkgs/main           
tensorflow                     2.9.1 mkl_py310h871f1f9_0  pkgs/main           
tensorflow                     2.9.1 mkl_py37h58a621a_0  pkgs/main           
tensorflow                     2.9.1 mkl_py38h96f9fba_0  pkgs/main           
tensorflow                     2.9.1 mkl_py39hb9fcb14_0  pkgs/main

と出力されるので、「gpu」から始まっているビルドを確認してインストールすれば、望んだ環境をインストールすることができます。「gpu_py xx」のxxの部分でPythonのバージョンも指定することができるみたいですね。

ただ、私の環境下ではTensorFlow=2.9.1のGPU環境は上手くインストールできませんでした。エラーメッセージを見て、解決できれば、また追記したいと思います。

これからの環境構築（メモ）

私が試した環境は以下の通りです。

Ubuntu 22.04.1 LTS
Intel(R) Core(TM) i7-7700K
Memory 16G
GeForce GTX 1080 Ti

これから新たしく機械学習の環境を整える際には

ubuntu-drivers devices

で、recommendされたnvidia-driverをインストールして、上記手順でTensorFlow・CUDA・cuDNNをインストールすれば、環境構築が終わるわけですね。
※「nouveau」を無効化する必要があるかもしれません。
※nvidia-driverはマシンに一度インストールすれば良いので、この章はマシンを再度セットアップし直さない限り、初回以外は参照する必要はありません。

謝辞

@hatsuyuki396様、有益な情報をありがとうございました。