Pytorch, torch-geometric 環境構築 (Linux)

概要

１ヶ月ほど前に研究室からRTX3090の計算機を貸していただきました。
またつい先日ですがpytorch公式からgeforce rtx3080~3090のグラボにも対応したとのことです。そのため実際に環境構築したので躓いたところなども含めメモ

pytorchは知っているが、geometricについてご存知でない方もいらっしゃるかもしれません。ざっくり言えば、
pytorch      : Neural Network 全般のライブラリ
geometric : Graph NN ($\in$ Neural Network)特化のライブラリ

ですので、geometricを使えば、様々なデータセットを、データセットの形式の違いに囚われずに、共通化したメソッドで簡単にミニバッジを作ったり、前処理をしたり、効率的なグラフ畳み込みを行えたりなど恩恵があります。

グラフ関連のNeural Networkは扱わないからpytorchだけでいいよ、という人は最後のgeometricのインストールだけ読み飛ばしてもらえればと思います

著者の環境

OS    :  Linux(ubuntu 20.04 LTS)
GPU : NVIDIA Geforce RTX3090

動作したバージョン構成

pytorch関連

python : 3.8.5 + pip3 (python3.9は非対応のため✖)
nvidia-driver : 455
cuda : 11.0(11.1でも良い)
cudnn : 8.0.3
pytorch : 1.7.0

torch-geometric関連

torch-cluster 1.5.8
torch-scatter 2.0.5
torch-sparse 0.6.8
torch-spline-conv 1.2.0
torch-geometric 1.6.1

環境構築詳細

それでは本題に行きましょう
長いので目次です

1. アンインストール
2. インストール(NVIDIA driver)
3. インストール(Cuda)
4. インストール(Cudnn)
5. インストール(pytorch)
6. インストール(torch-geometric)

1. アンインストール

すでに違うバージョンで環境構築をしていて、動作しないトラブルを抱えている方は一旦それらをアンインストールしましょう

nvidiaドライバ、cuda、cudnnのアンインストール

$ sudo apt-get --purge remove cuda*
$ sudo apt-get --purge remove nvidia-cuda-*
$ sudo apt-get --purge remove libcuda*
$ sudo apt autoremove
$ reboot

pytorch, torch-geometric のアンインストール

sudo pip3 uninstall torch
sudo pip3 uninstall torchvision
sudo pip3 uninstall geometric*

2. インストール(NVIDIA driver)

ドライバを入れましょう。ここで注意ですが、NVIDIAのグラボの場合、デフォルトのnouveauというドライバと競合してしまいます。これを避けるためにこいつを無効化します。

lsmod | grep nouveau #ここで何か表示されたら下のコマンド
blacklist nouveau
options nouveau modeset=0
sudo update-initramfs -u\
sudo reboot #再起後、ディスプレイの解像度が下がれば無効化成功

ではグラボのドライバのインストールに行きましょう

ubuntu-drivers devices #ここで対応のバージョンが推奨される
sudo add-apt-repository ppa:graphics-drivers/ppa
sudo apt update
apt-cache search nvidia | grep nvidia-driver-455
sudo apt-get install nvidia-driver-455
sudo reboot

自分の環境ではバージョン455が推奨されたため上のコマンドになりました。違うバージョンを推奨された場合は適宜変えてください

3. インストール(Cuda)

今回はRTX3080~に対応した環境構築が主題のため以下のcuda11.0のインストールサイトからインストールします
https://developer.nvidia.com/cuda-11.0-download-archive
上のページで環境を聞かれるので(Linux > x86_64 > Ubuntu > 20.04 > deb(network))と入力するとインストールのためのコマンドが表示されるのでそのとおりに叩く。これでインストールは済むはずです。

後はパスの設定をします。以下のパス設定のコマンドを '~/.bashrc'に書き込んでsourceして反映しましょう

export PATH=/usr/local/cuda-11/bin:${PATH}
export LD_LIBRARY_PATH=/usr/local/cuda-11/lib64:${LD_LIBRARY_PATH}

正しくパスが通っていればcudaに含まれているnvcc(コンパイラ)のコマンドが動作するはずです

nvcc -V #正しくバージョンが表示されたらok
nvidia-smi #正しくバージョンが表示されたらok

4. インストール(Cudnn)

インストールしたCudaのバージョンに対応したcuDNNをインストールしましょう。対応しているかどうかは下のように公式を見れば明らかでわかりやすいです(今回はversion8.0.3)

公式のページからログインが済めば、上のスクショのように、自分のCudaと対応したバージョンのcuDNNをクリックしてください
するとDebファイル一覧が表示されるため、
・cuDNN Runtime Library for Ubuntu * x86_64 [Deb]
・cuDNN Developer Library for Ubuntu * x86_64 [Deb]
・cuDNN Code Samples for Ubuntu * x86_64 [Deb]
の３つを選択してダウンロードしましょう

ダウンロードしたフォルダまでcdして、cuDNNをインストールしましょう(注意:インストールの順番は以下のとおりにしてください、依存関係があります)

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

5. インストール(Pytorch)

注意:冒頭にも書いていますがPythonのバージョンは3,3以上3.8以下にしてください(何事も3.7が一番無難かも)。
apt-get install pythonでバージョン指定せずにPythonインストールするとバージョン3.9入れられた気がするので気をつけてください

自分はpipでライブラリを管理してもらってるためpipで入れます。公式のページ (https://pytorch.org/get-started/locally/) へ行き、(Stable(1.7.0) > Linux > Pip > Python > 11.0)と入力したらこのコマンドを走らせて、と言われるのでそのとおりに叩きます

しかしもし何かしらエラーを言われたら
pip install ~ を sudo pip3 install に修正してください

正しくpytorchがインストールできたか、またpytorchがグラボを使用できているか確認しましょう

python3
>>> import torch #何もエラーが出なければok
>>> print(torch.__version__) #バージョンが出ればok
>>> torch.cuda.is_available() #Trueと出ればok

6. インストール(torch-geometric)

公式 (https://pytorch-geometric.readthedocs.io/en/latest/notes/installation.html) の言う通りすればできました。

pip install torch-scatter==latest+${CUDA} -f https://pytorch-geometric.com/whl/torch-${TORCH}.html
pip install torch-sparse==latest+${CUDA} -f https://pytorch-geometric.com/whl/torch-${TORCH}.html
pip install torch-cluster==latest+${CUDA} -f https://pytorch-geometric.com/whl/torch-${TORCH}.html
pip install torch-spline-conv==latest+${CUDA} -f https://pytorch-geometric.com/whl/torch-${TORCH}.html
pip install torch-geometric

\${CUDA}は(cpu, cu92, cu101, cu102, cu110) のいずれか自分の環境に合ったものを、\${TORCH}は(1.4.0, 1.5.0, 1.6.0, 1.7.0)のいずれか合ったものを。

動作確認

python3
>>> from torch_geometric.data import Data
>>> edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]], dtype=torch.long)
>>> x = torch.tensor([[-1], [0], [1]], dtype=torch.float)
>>> data = Data(x=x, edge_index=edge_index)

#>>> Data(x=[3, 1], edge_index=[2, 4])とでれば成功です

お疲れ様でした。