JetsonNano2GBにTensorFlow環境を構築する方法

えびかずき

こんにちは、えびかずきです。

今回は、JetsonNano2GBにTensorFlowをインストールして動かすまでの手順を説明していきます。

こんな人におすすめ：
・JetsonNano2GBにTensorFlowを入れて、推論用のエッジAIとして活用したい。
・JetsonNano2GBを買ったはいいが、使い方がよくわからないので放置している。

Tensorflowの環境構築については、Nvidia公式のDEEP LEARNING FRAMRWORKS DOCUMENTATIONに記載があるので基本的にはそれに従って作業を進めることで環境を構築できます。

ただし、いくつか気を付けるべきポイントがあるので一つ一つ順を追って説明していきます。

1. 開発環境
2. 事前準備
3. TensorFlowのインストール
4. TensorFlowの動作テスト
- 4.1. ソースコード
- 4.2. テストの実行
5. まとめ
6. 参考URL

開発環境

・Jetson Nano 2GB
・JetPack 4.5.1
・Python 3.6.9
・TensorFlow 2.3.1
・Pillow,Numpy

事前準備

１．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

２．pipのインストール

初期状態だとpipが使えません。

下のようにしてpipをインストールしておきます。

$ sudo apt-get install python3-pip
$ sudo pip3 install -U pip testresources setuptools==49.6.0

３．必要なPythonパッケージのインストール

Tensorflowを動かすのに必要なPythonパッケージを事前にインストールしておきます。

下のようにpip3でインストールします。

$ sudo pip3 install -U numpy==1.19.4 future==0.18.2 mock==3.0.5 h5py==2.10.0 keras_preprocessing==1.1.1 keras_applications==1.0.8 gast==0.2.2 futures protobuf pybind11

※ここは結構時間がかかります。特にNumpyのインストールに時間を喰います。私の場合は15分程度かかりました。

TensorFlowのインストール

下のコマンドでTensorFlowのバージョンを指定してインストールすることができます。

$ sudo pip3 install --extra-index-url https://developer.download.nvidia.com/compute/redist/jp/v45 tensorflow==2.3.1+nv20.12

ちなみに上のコマンドでJetPack4.5にインストールできるTensorfloｗのバージョンは、以下のリンク先にあるバージョンに限られます。

https://developer.download.nvidia.com/compute/redist/jp/v45/tensorflow/

この中から自分の使いたいバージョンを選んでインストールしましょう。

さてこれで環境構築は完了です。

注意：
公式のドキュメントでは下のように「==」と「+」の両側にスペースが入っていますが、これだとうまくいきません。スペースは不要です。

“tensorflow == 2.3.1 + nv20.12”

TensorFlowの動作テスト

TensorFlowの環境構築ができたので、動作テストをやってみましょう。

今回は、過去記事のディープラーニングで画風変換を実装する方法で紹介した画風変換の学習済みモデルを使って、TensorFlowの動作テストをやってみます。

今回は下の左側のゴッホの絵のスタイルを、右の太陽の塔の写真に適用するというディープラーニングモデルを使用します。

詳しい技術説明については過去記事をご参照ください。

ソースコード

実装したソースコードは以下のとおり。

# test.py

from tensorflow.python.keras.models import load_model
from tensorflow.python.keras.preprocessing.image import load_img, img_to_array, array_to_img
import numpy as np

#modelへ保存データを読み込み
model_gen = load_model('test10.h5')

# 画風変換の実行
input_shape = (224, 224, 3)
input_size = input_shape[:2]
img_test = load_img(
    'img/test/taiyou.jpg', 
    target_size=input_size
)
img_arr_test = img_to_array(img_test)
img_arr_test = np.expand_dims(
    img_to_array(img_test), 
    axis=0
)

pred = model_gen.predict(img_arr_test)

# 変換後の画像の表示
img_pred = array_to_img(pred.squeeze())
img_pred.save("test.jpg")

ここでは既に学習済みのモデルとして「test10.h5」というh5ファイルを使ってます。

このファイルが欲しい人は過去記事のGPUを使うとどれだけディープラーニングが速くなるか？に添付したGithubリポジトリから取得できますのでどうぞ。