RTX3090+Ubuntu20.04でPyTorchを動かす

RTX3090 + Ubuntu20.04でPyTorch-1.7.0が動いたので、メモ．

ドライバのインストール

$ sudo add-apt-repository ppa:graphics-drivers/ppa
$ sudo apt update
$ sudo apt install nvidia-driver-455

インストール後に再起動

$ sudo reboot

CUDA-11.1のインストール

RTX3000シリーズはCUDA11にしか対応してない?

CUDA Toolkit 11.1 Update 1 Downloads | NVIDIA Developerの通りにコマンドを実行する

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

cuDNN-8.0.5のインストール

NVIDIA cuDNN | NVIDIA DeveloperのDownload cuDNNからダウンロードしてきたlibcudnn8-dev_8.0.5.39-1+cuda11.1_amd64.debをインストールする。

$ sudo apt install -y ./libcudnn8-dev_8.0.5.39-1+cuda11.1_amd64.deb

PyTorchのインストール

PyTorchのGet Startedページで

• Stable (1.7.0)
• Linux
• Pip
• Python
• 11.0

を選んで出てきたコマンドを実行する

$ pip install torch==1.7.0+cu110 torchvision==0.8.1+cu110 torchaudio===0.7.0 -f https://download.pytorch.org/whl/torch_stable.html

---

これでRTX3090+Ubuntu20.04+PyTorch1.7.0が動いた。