Tensorflow 2.x & PyTorch 1.6でGPUを使う【Ubuntu 20.04】

更新

2020.09.05 pipenvでPyTorchのインストール方法を変更した

この記事が目指す環境

・Ubuntu 20.04 LTS
・pyenv (Pythonのバージョン管理ツール)
・Python 3.x (3.6以上推奨)
・pipenv (仮想環境の作成等)
・Tensorflow 2.x (GPU対応)
・PyTorch 1.6.x (GPU対応)
・CUDA 10.1
・CuDNN 7.x

Pythonのインストールについては以下の記事に書いてあります。

GPUドライバーのインストール

「ソフトウェアとアップデート」のアプリケーションを起動します。

「追加のドライバー」から以下の画像のように、オープンソースでないnvidia-driver-440のGPUドライバーを選択し、再起動します。

CUDAツールキットのインストール

以下のコマンドを端末で実行しましょう。

sudo apt install nvidia-cuda-toolkit

インストールが無事に終了すると「nvcc -V」コマンドが使用できるようになります。

cuDNN SDKのインストール

NVIDIA Developerのサイトに行きアカウントを作成してログインします。NVIDIA cuDNNのページに行き、「Download cuDNN」をクリックします。
ソフトウェアライセンス契約の条件に同意すると以下のような項目が出てくるので、一番下にある「Archived cuDNN Releases」を選択します。

過去にリリースされたcuDNN一覧が表示されるので、「Download cuDNN v7.6.5(November 18th, 2019), for CUDA 10.1」をクリックし、「cuDNN Library for Linux」を選択してダウンロードします。

ダウンロードしたファイルを解凍して、解凍したcuDNNのディレクトリ内に端末で移動します。

移動したら以下の3つのコマンドを実行します。

sudo cp cuda/include/cudnn.h /usr/lib/cuda/include/
sudo cp cuda/lib64/libcudnn* /usr/lib/cuda/lib64/
sudo chmod a+r /usr/lib/cuda/include/cudnn.h /usr/lib/cuda/lib64/libcudnn*

次にシェルの環境変数を設定します。bashを使用している場合は「~/.bashrc」に、zshを使用している場合は「~/.zshrc」に以下の2行を追記して保存してください。

export LD_LIBRARY_PATH="/usr/lib/cuda/lib64:$LD_LIBRARY_PATH"
export LD_LIBRARY_PATH="/usr/lib/cuda/include:$LD_LIBRARY_PATH"

保存できたら、以下のコマンドで変更を適用します。

# bashの場合
source ~/.bashrc

# zshの場合
source ~/.zshrc

Tensorflow 2.xのインストール

念の為一度再起動したら以下のコマンドを実行しましょう。

# pipの場合
pip install tensorflow

# pipenvの場合
pipenv install tensorflow

既にインストールしていた場合は一度アンインストールを行いましょう。

# pipの場合
pip uninstall tensorflow
pip install tensorflow

# pipenvの場合
pipenv uninstall tensorflow
pipenv install tensorflow

GPUドライバーが正しく認識できているか確認しましょう。

# 任意の方法で以下のPythonコードを実行する。
import tensorflow as tf
print(tf.config.list_physical_devices('GPU'))

最後の行に以下のように表示されていればOKです。

[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]

このままでGPUの認識はしますが、一部のGPUではmodel.fit()等を実行した際にエラーを吐くことがあるので、以下のように書くことをおすすめします。

import tensorflow as tf
gpu_devices = tf.config.experimental.list_physical_devices('GPU')
for gpu_device in gpu_devices:
   tf.config.experimental.set_memory_growth(gpu_device, True)

PyTorch 1.6.xのインストール

インストールする際のコマンドはPyTorch公式サイトにあります。

pipでインストールする際は以下のコマンドを実行してインストールします。

# pipでインストールする場合
pip install torch==1.6.0+cu101 torchvision==0.7.0+cu101 -f https://download.pytorch.org/whl/torch_stable.html

pipをpipenvに書き換えると以下のようなエラーが出てしまいます。

上記の問題を回避する為にtorchとtorchvisionはpipenvの仮想環境内でpipインストールしましょう。

# pipenvでインストールする場合
pipenv shell
pipenv install future numpy pillow
pip install torch==1.6.0+cu101 torchvision==0.7.0+cu101 -f https://download.pytorch.org/whl/torch_stable.html

GPUドライバーが正しく認識できているか確認しましょう。

# 任意の方法で以下のPythonコードを実行する。
import torch
print(torch.cuda.is_available())

Trueと表示されれば問題なくインストールされています。