【jetson XAVIER NX】tensorflowでGPUを利用できる環境を構築

以下環境で動作確認を行いました。

jetson XAVIER NXにtensorflowの環境を構築する手順は4ステップになります。

環境の確認

Jetpackのバージョンを確認

まずインストールしたOS(jetpack)を確認します。

sudo apt-cache show nvidia-jetpack

2020年9月1日時点で最新のバージョンとなっています。

#Package: nvidia-jetpack
Version: 4.4-b144
Architecture: arm64
Maintainer: NVIDIA Corporation

Python3のバージョンを確認

python3のバージョン確認コマンドを打ちます。

python3 --version

OSにデフォルトでインストールされていたバージョンが3.6.9でした。

Python 3.6.9

pythonの仮想環境の構築

pipのインストール

Virtualenvをインストールするためにpip3をaptでインストールします。

sudo apt install python3-pip

Virtualenv のインストール

pythonやtensorflowのバージョンを簡単に変えられるように、今回はVirtualenvの仮想環境上に構築します。

sudo pip3 install virtualenv

Virtualenvの環境構築

virtualenvでpythonの仮想環境を作成します。

virtualenv tf_env

tf_envは仮想環境名なので、任意の文字列を使ってください。

Virtualenv 仮想環境の有効化

virtualenvの環境を起動します。

source tf_env/bin/activate

tf_envは仮想環境名なので、先ほど作った環境名を使ってください。

tensorflowのインストール

nvidia公式の手順に則って実施していきます。
https://docs.nvidia.com/deeplearning/frameworks/install-tf-jetson-platform/index.html

TensorFlowに必要なシステムパッケージをインストール

sudo apt-get update 
sudo apt-get install libhdf5-serial-dev hdf5-tools libhdf5-dev zlib1g-dev zip libjpeg8-dev liblapack-dev libblas-dev gfortran

pip3のインストールとアップグレード

sudo apt-get install python3-pip 
sudo pip3 install -U pip testresources setuptools

Pythonパッケージの依存関係をインストール

sudo pip3 install -U numpy==1.16.1 future==0.17.1 mock==3.0.5 h5py==2.9.0 keras_preprocessing==1.0.5 keras_applications==1.0.8 gast==0.2.2 futures protobuf pybind11

TensorFlowのインストール

tensorflowのバージョン2.1.0を指定してインストールします。

pip3 install --extra-index-url https://developer.download.nvidia.com/compute/redist/jp/v44 tensorflow==2.1.0+nv20.04

sudo pip3 install --pre --extra-index-url https://developer.download.nvidia.com/compute/redist/jp/v44 tensorflow

tensorflow動作確認

python3
>>>import tensorflow

エラーにならずにインポートできればOK

>>> tensorflow.__version__

バージョンの確認ができます。

GPUの動作確認

GPUの確認

pythonで以下のコードを実行すると認識できているデバイスが表示されます。

>>>from tensorflow.python.client import device_lib
>>>device_lib.list_local_devices()

GPUが認識されていることがわかりました。

incarnation: 8185341057002331750
physical_device_desc: “device: XLA_GPU device”
, name: “/device:GPU:0”
device_type: “GPU”
memory_limit: 1905582080
locality {
bus_id: 1
links {
}
}

mnistの実行

CPUとGPUそれぞれでmnistデータの学習を行ってみました。

import tensorflow as tf import os 

#GPUをオフにする場合 #os.environ["CUDA_VISIBLE_DEVICES"]="-1" 

#認識されているデバイスを確認 
from tensorflow.python.client import device_lib print(device_lib.list_local_devices()) 
# CPUでmnistを実行 
with tf.device('/cpu:0'): 
    mnist = tf.keras.datasets.mnist (x_train, y_train), 
    (x_test, y_test) = mnist.load_data() 
    x_train, x_test = x_train / 255.0, x_test / 255.0 
    model = tf.keras.models.Sequential([ 
    tf.keras.layers.Flatten(), 
    tf.keras.layers.Dense(512, activation=tf.nn.relu), 
    tf.keras.layers.Dropout(0.2), 
    tf.keras.layers.Dense(10,activation=tf.nn.softmax) 
    ]) 
    model.compile(optimizer='adam', 
    loss='sparse_categorical_crossentropy', 
    metrics=['accuracy']) 
    model.fit(x_train, y_train, batch_size=100, epochs=5) 
    model.evaluate(x_test, y_test) 
# GPUでmnistを実行 
with tf.device('/device:GPU:0'): 
    mnist = tf.keras.datasets.mnist (x_train, y_train), 
    (x_test, y_test) = mnist.load_data() 
    x_train, x_test = x_train / 255.0, x_test / 255.0 
    model = tf.keras.models.Sequential([ 
    tf.keras.layers.Flatten(), 
    tf.keras.layers.Dense(512, activation=tf.nn.relu), 
    tf.keras.layers.Dropout(0.2), 
    tf.keras.layers.Dense(10, activation=tf.nn.softmax) 
    ]) 
    model.compile(optimizer='adam', 
    loss='sparse_categorical_crossentropy', 
    metrics=['accuracy']) 
    model.fit(x_train, y_train, batch_size=100, epochs=5) 
    model.evaluate(x_test, y_test)

実行中にリソースを確認すると、それぞれの処理のタイミングでCPUとGPUの使用率が高くなることを確認できました。

ロボくん

tensorflow2.1ではデバイスを指定しないとデフォルトでGPU処理される仕様になっているよ。

以上でtensorflowのインストールと動作確認は完了になります。

まとめ

jetson XAVIER NXにtensorflowをインストールする方法を紹介しました。

GPUが本当に使えているかの動作確認もmnistを使った解析で確認することができたと思います。

公式のドキュメント通りにやっても失敗する部分があったので、気を付けて実施してみてください。