概要

PyTorch 1.7にてRTX 3090に対応してくださったので、その環境構築のメモとして残しておきます。

目標の環境

ハード
- GPU : RTX 3090 ← 高い
ソフト
- OS : Ubuntu 18.04
- CUDA : 11.0
- cuDNN : 8.0.4
- PyTorch : 1.7.1

これで torch.cuda.is_available() がTrueになって欲しい…！

手順

0. GPUを差し替える

(この作業がない人は飛ばしてください)

もともとGTX1080Ti を使用していたので、まずはハードを差し替えました。
特にドライバ関係をアンインストールすること無く、差し替えてしまいました。(大丈夫だったのかな...)
一応、GTX1080Ti 時の環境を以下に示します。

OS : Ubuntu 18.04
CUDA : 10.2
cuDNN : 8.0.3
PyTorch : 1.6.0

そして、起動したら解像度がおかしくなってました。
そりゃ、もともとのドライバ関係でおかしくなるのも当然だと思います。

1. 既存のドライバ、CUDA、cuDNN、PyTorch等の削除

(この作業が必要ない人は飛ばしてください)

ターミナルで以下のコマンドを使用して、残っている nvidia driver、CUDA、cuDNNを消し去ります。

$ sudo apt remove --purge nvidia*
$ sudo apt remove --purge cuda*
$ sudo apt remove --purge libcudnn*
$ sudo apt remove --purge libnvidia*
$ sudo apt autoremove

PyTorch等を消し去ります。僕はpipで管理してましたので、以下のコマンドを打ちました。

$ pip uninstall torch
% pip uninstall torchvision

念の為、ここでPCに再起動をかけてやりましょう。

$ sudo reboot

2. nvidia-driverをインストール

以下のコマンドを使用して、RTX3090に合ったnvidia-driverを確認してください。

$ sudo ubuntu-drivers devices

僕は、nvidia-driver-455がrecommendedされましたので、以下のコマンドを打ちインストールしました。

$ sudo apt install nvidia-driver-455 # 2021/01/07 時点

インストールを終えたら、ここで再起動しましょう。

$ sudo reboot

たぶん、解像度が元通りになってると思います。nvidia-smiでドライバの状態等々を確認しましょう。

$ nvidia-smi

3. CUDAをインストール

nvidia-smiで確認したところ、右上にCUDA 11.1って書いてある。でも、PyTorch 1.7のサポートはCUDA 11.0。

僕もよくわかってないのですが、CUDA 11.0を入れてしまいましょう...。
CUDA 11.0のアーカイブサイトからLinux > x86_64 > Ubuntu > 18.04 > deb(network) を選びました。そこに記載されているコマンドを入力してください。

### 以下上記サイトのコピペをしてください。 ###
$ wget https://developer.download.nvidia.com/....
...
$ sudo apt-get update
### ここまで ###

ここで最後の行のコマンドに注意が必要です。
以下のようにバージョンを指定しないと、最新のバージョンがインストールされてしまうので、cuda-11-0としておきましょう。

$ sudo apt-get -y install cuda-11-0

パスを通して無い方は、~/.bashrcの最終行に以下を追加しましょう。

$ sudo vi ~/.bashrc

~/.bashrc

## CUDA
export PATH="/usr/local/cuda/bin:$PATH"
export LD_LIBRARY_PATH="/usr/local/cuda/lib64:$LD_LIBRARY_PATH"

vimでは、[Ctrl]+[Shift]+[v]でペーストして、[Esc]キーでNORMALモードにして、:wqで上書き保存してください。
これでCUDA 11.0がインストールできました。以下のコマンドで確認してみましょう。

$ nvcc -V

4. cuDNNをインストール

cuDNNのアーカイブサイトから最新のバージョン for CUDA 11.0を選びましょう。僕は、v8.0.4 を選びました。
その中の、

cuDNN Runtime Library for Ubuntu18.04 x86_64 (Deb)
cuDNN Developer Library for Ubuntu18.04 x86_64 (Deb)
cuDNN Code Samples and User Guide for Ubuntu18.04 x86_64 (Deb)

を選び、ダウンロードしました。

僕は、~/Downloadsディレクトリにダウンロードされるので、

$ cd Downloads

してから、以下の順番でコマンドを打ってcuDNNをインストールしましょう。

$ sudo dpkg -i libcudnn8_8.0.4.30-1+cuda11.0_amd64.deb
$ sudo dpkg -i libcudnn8-dev_8.0.4.30-1+cuda11.0_amd64.deb
$ sudo dpkg -i libcudnn8-samples_8.0.4.30-1+cuda11.0_amd64.deb

念の為、ここでPCに再起動をかけてやりましょう。

$ sudo reboot

5. PyTorchをインストール

PyTorchの公式サイト通りのコマンドを打ってインストールしましょう。2021/01/07 時点で以下のコマンドでした。

$ pip install torch==1.7.1+cu110 torchvision==0.8.2+cu110 torchaudio===0.7.2 -f https://download.pytorch.org/whl/torch_stable.html

結果

PyTorchでの確認

$ python
Python 3.6.7 (default, Sep  7 2020, 17:00:49) 
[GCC 7.5.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.cuda.is_available()
True

CUDAの確認

$ nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2020 NVIDIA Corporation
Built on Wed_Jul_22_19:09:09_PDT_2020
Cuda compilation tools, release 11.0, V11.0.221
Build cuda_11.0_bu.TC445_37.28845127_0

nvidia-smiでの確認

$ nvidia-smi
Thu Jan  7 23:35:29 2021       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 460.27.04    Driver Version: 460.27.04    CUDA Version: 11.2     |
|-------------------------------+----------------------+----------------------+
| 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 3090    On   | 00000000:65:00.0  On |                  N/A |
| 51%   67C    P2   295W / 350W |   5855MiB / 24267MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

なんでCUDA 11.2 になってるんだろう。

あとがき

この後PyTorchでモデルを学習させましたが、しっかり動いてました。
乱雑に書いてしまったところもあるので、もし不明な部分や間違いがありましたらコメント等ください m(_ _)m

便利な世の中ですね。