GPUで機械学習環境構築するのって色んなドライバーのバージョンやらなんやらで禿げ上がりそうになるんですが、色々やってみてこれが多分一番楽そうなのでご紹介します。

想定環境

・Linux Mint 19.3
・RTX2070 SUPER
・tensorflowを使用
・Anacondaを使用

ドライバーを入れる

まずはNVIDIAのドライバーを入れます。
nouveauが動いている人は無効化する必要があるそうですが、私の環境では動いてなかったので特に設定していません。

sudo add-apt-repository ppa:graphics-drivers/ppa
sudo apt update
ubuntu-drivers devices

これで推奨のドライバー等が出てきますので、素直に最新のものを入れました。

modalias : pci:v000010DEd00001E84sv00001462sd000037B3bc03sc00i00
vendor   : NVIDIA Corporation
driver   : nvidia-driver-455 - third-party free recommended
driver   : nvidia-driver-440-server - distro non-free
driver   : nvidia-driver-450-server - distro non-free
driver   : nvidia-driver-435 - distro non-free
driver   : xserver-xorg-video-nouveau - distro free builtin

sudo apt install nvidia-driver-455

インストールできると再起動して以下を実行して、ドライバーの情報が表示されれば成功です。

nvidia-smi
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 455.32.00    Driver Version: 455.32.00    CUDA Version: 11.1     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  GeForce RTX 207...  On   | 00000000:01:00.0 Off |                  N/A |
| 31%   32C    P8     6W / 215W |     73MiB /  7982MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      1158      G   /usr/lib/xorg/Xorg                 71MiB |
+-----------------------------------------------------------------------------+

conda仮想環境を作る

今回はconda installにtensorflow, cuda, cudnnのバージョン合わせをお任せしようというやり方です。
なので、まずconda環境を作ります。

conda create --name NAME python=3.7

環境ができているかどうかは

conda info -e

で確認できます。
無事に作成できていたら、

conda activate NAME

で仮想環境をアクティベートしましょう。

tensorflow, cuda, cudnnインストール

ドライバーのバージョンが古い場合のみ対応するcuda→tensorflowという順番で、tensorflowのバージョンを特定してください。
例えばドライバーが384.81しか入らなかったら、cuda9.0になって、それに対応するtensorflowはtensorflow_gpu-1.12.0です。下記リンクを参照ください。

私の場合は455.32を入れれたので、cuda、tensorflowも最新のものを利用できます。以下を実行します。

conda install tensorflow-gpu

これでtensorflowだけでなく、cuda、cudnnも適切なバージョンが入ってしまいます。

# conda list tensorflow
tensorflow                2.2.0           gpu_py37h1a511ff_0
tensorflow-base           2.2.0           gpu_py37h8a81be8_0
tensorflow-estimator      2.2.0              pyh208ff02_0
tensorflow-gpu            2.2.0                h0d30ee6_0

# conda list cuda
cudatoolkit               10.1.243             h6bb024c_0

# conda list cudnn
cudnn                     7.6.5                cuda10.1_0

以下を実行してエラーが特になければGPUを正常に認識しています。

tensorflow.config.list_physical_devices("GPU")

まとめ

GPU計算環境を構築するのにすごい時間を奪われたので、最終的に簡単にできた方法を共有します。