🎃

Ubuntu 22.04でcudaインストールしPyTorchとTensorFlowをGPUで動かす方法

2023/08/11に公開

TensorFlow

Ubuntu 22.04でPyTorchとTensorFlowをGPUで動かすためには、以下の手順を実行することでセットアップを行うことができます。ただし、公式のサポート状況やパッケージの互換性は変更されることがあるため、公式サイトのドキュメントやリリースノートを確認することを推奨します。

1. NVIDIA Driverのインストール

まず、NVIDIAのGPUドライバをインストールします。

以下を実行すると、おすすめのNVIDIA Driverのバージョンを教えてくれます。

ubuntu-drivers devices

私の環境では今の出力となります。

vendor   : NVIDIA Corporation
model    : GA102 [GeForce RTX 3090]
driver   : nvidia-driver-525-open - distro non-free
driver   : nvidia-driver-535-server - distro non-free
driver   : nvidia-driver-525 - distro non-free
driver   : nvidia-driver-470-server - distro non-free
driver   : nvidia-driver-535-open - distro non-free
driver   : nvidia-driver-525-server - distro non-free
driver   : nvidia-driver-535-server-open - distro non-free
driver   : nvidia-driver-470 - distro non-free
driver   : nvidia-driver-535 - distro non-free recommended
driver   : xserver-xorg-video-nouveau - distro free builtin

この中でrecommendedのnvidia-driver-535をインストールすることにします。- distro non-freeなどの文字は無視してください。

sudo apt update
sudo apt install nvidia-driver-535

再起動して、正常にインストールされたことを確認します。

nvidia-smi

以下のように表示されればインストール成功です。

2. CUDA Toolkitのインストール

↑このページを見ると、NVIDIA Driver Version: 535.86.05ならば現時点（2023年8月）現在で最新のCUDA 12.2 Update 1にも対応していることがわかります。

公式のCUDA Toolkitのダウンロードページから、Ubuntu 22.04に対応するCUDAバージョンをダウンロードしてインストールします。PyTorch、Tensorflowを動かす時にはモデルが新すぎると動かないコードがたくさんあるため、なるべく低いバージョンのToolkitを選ぶことにします。今回は11.7を選びました。以下のコマンドでインストールできます。

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb
sudo dpkg -i cuda-keyring*.deb
sudo apt-get update
sudo apt-get install cuda-11-7

3. 環境変数の設定

.bashrcや.zshrcなどのシェル設定ファイルに、CUDA関連の環境変数を追加します。ここで、/usr/local/cuda-11-7のようにバージョン番号を入れることを忘れないようにしてください。

# Set PATHS for CUDA
export PATH=/usr/local/cuda-11.7/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-11.7/lib64:$LD_LIBRARY_PATH

fish shellの場合は、~/.config/fish/config.fishに

# Set LD_LIBRARY_PATH for CUDA
set -x BNB_CUDA_VERSION 117
set -x LD_LIBRARY_PATH /usr/local/cuda-11.7/lib64

を追記しましょう。

4. cuDNNのインストール

NVIDIAの公式サイトから、ダウンロードページにアクセスして、使用しているCUDAバージョンに対応するcuDNNをダウンロードしてインストールします。今回はcudnn 8.5.0をインストールしたいと思います。以下のページから.debのファイルをダウンロードできます。ここではメールアドレスの登録が必要になります。

ダウンロードしたら、以下のコマンドでインストールします。

sudo dpkg -i cudnn-local-repo-ubuntu2204-8.5.0.96_1.0-1_amd64.deb
sudo cp /var/cudnn-local-repo-ubuntu2204-8.5.0.96/cudnn-local-7ED72349-keyring.gpg /usr/share/keyrings/
sudo apt update
sudo apt-get install libcudnn8=8.5.0.96-1+cuda11.7
sudo apt-get install libcudnn8-dev=8.5.0.96-1+cuda11.7

5. PyTorchのインストール

公式のPyTorchサイトから、最新のインストールコマンドを取得して実行します。CUDAバージョンに合わせて適切なコマンドを選択します。

pip install torch torchvision

6. TensorFlowのインストール

公式のTensorFlowサイトから、最新のインストールコマンドを取得して実行します。CUDAバージョンに合わせて適切なコマンドを選択します。

pip install tensorflow

上記の手順で、Ubuntu 22.04上でPyTorchとTensorFlowをGPUで動かすための環境をセットアップすることができます。ただし、各パッケージのバージョンや互換性に関する最新の情報は、公式のドキュメントやフォーラムで確認することを推奨します。