Ampere世代GPUでtensorflow-2.4.0 + cuda11.0 + cudnn8 環境を構築したら苦労した

はじめに

Ampere世代のGPU　RTX3070を購入してtensorflow2.4.0の機械学習環境を構築しましたが、かなり難渋しました。
理由はtensorflow2.4.0がCuda11.0に対応しているが、Ampere世代のGPUはcuda11.1で動作するようになっていたためです。

以下の参考手順は一例でありますの参考程度かつ自己責任でお願いいたします。
またエンジニアではありませんので一部読みにくい表記などがございますが、ご容赦ください。

環境など
CPU: Ryzen9 5950x
GPU: RTX3070(Ampere世代）
OS: Ubuntu20.04

参考サイト
異なるバージョンのCUDAを使い分ける単純な方法

nvidia-driverのインストール

下記のとおりに、version 460がinstallされました（20210122時点）

bash

sudo add-apt-repository ppa:graphics-drivers/ppa
sudo apt update
ubuntu-drivers devices
sudo apt install nvidia-driver-<version>

cuda11.0のインストール

NVIDIAのサイトでcuda11.0をインストールします。

bash

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
wget http://developer.download.nvidia.com/compute/cuda/11.0.2/local_installers/cuda-repo-ubuntu2004-11-0-local_11.0.2-450.51.05-1_amd64.deb
sudo dpkg -i cuda-repo-ubuntu2004-11-0-local_11.0.2-450.51.05-1_amd64.deb
sudo apt-key add /var/cuda-repo-ubuntu2004-11-0-local/7fa2af80.pub
sudo apt-get update
sudo apt-get -y install cuda-toolkit-11-0

パスを通す

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

cudnn8のインストール

cudnnはNvidiaのデベロッパーサイトから会員登録してダウンロードします。
対応するOSなどを参考にしてダウンロード、インストールします。

bash

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

cudnnsampleを動かしてみる。

テストして動作するか確認のみ

bash

cp -r /usr/src/cudnn_samples_v8/ $HOME
cd $HOME/cudnn_samples_v8/mnistCUDNN
make clean && make
./mnistCUDNN

Test passed!　とでれば大丈夫そう！

tensorflow2.4.0のインストールするがsm_86 errorがでる。。。

さてここからが本題tensorflow2.4.0をインストールします。

pip3 install tensorflow==2.4.0

これでinstallできますが、機械学習などを実行すると下記errorがでてきます。
(学習、推論は可能だがエラーが頻出してなかなか実行されない）

Value 'sm_86' is not defined for option 'gpu-name'. 

これはAmpere世代のGPUはCUDA11.1以上が必要なのに対し、tensorflow2.4.0がCUDA11.0を前提に作られているためと思われます。

正しいかわかりませんが /usr/local/cuda-11.0/bin/ptxas　をcuda-11.1の　/usr/local/cuda-11.1/bin/ptxas　からコピーして入れ替えるとエラーなく動作します。

まずはcuda11.1をインストールしないことには/usr/local/cuda-11.1/bin/ptxasが入手できません。

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
wget https://developer.download.nvidia.com/compute/cuda/11.1.1/local_installers/cuda-repo-ubuntu2004-11-1-local_11.1.1-455.32.00-1_amd64.deb
sudo dpkg -i cuda-repo-ubuntu2004-11-1-local_11.1.1-455.32.00-1_amd64.deb
sudo apt-key add /var/cuda-repo-ubuntu2004-11-1-local/7fa2af80.pub
sudo apt-get updatesudo apt-get -y install cuda-toolkit-11-1

cuda11.1は使用しませんが、複数のcudaを入れておいても問題ないようです。
参考サイト

cuda11.0に再度パスを通しました。

bash

export PATH="/usr/local/cuda-11.0/bin:$PATH"
export LD_LIBRARY_PATH="/usr/local/cuda-11.0/lib64:$LD_LIBRARY_PATH"

さてptxasファイルをcuda11.1からcuda11.0にコピーします。もちろんファイルのバックアップは自分でしておきましょう。

bash

sudo cp /usr/local/cuda-11.1/bin/ptxas /usr/local/cuda-11.0/bin/

これでエラー（Value 'sm_86' is not defined for option 'gpu-name'. ）がでなくなり、推論結果自体は問題ないことを確認しました。

おわりに

Ampere世代GPUにtensorflow2.4.0の学習環境をつくりました。
Ampere世代がcuda11.1以降しか対応していないため面倒でした。
tensorflow2.4.0がcudaのいずれのバージョンで動いてくれるとありがたいのですがね。。。