[Amazon SageMaker] 最新のJetPack4.5でDLR (Deep Learning Runtime) 1.8.0を利用するための環境構築

1 はじめに

CX事業本部の平内（SIN）です。

エッジで機械学習の推論を行う場合、SageMaker Neo（モデル最適化）とJetson Nanoという組み合わせは、非常にコスパが高いと感じています。

今回は、2021/02/20 現在で最新のJetPack 4.5と、DLR(Deep Learning Runtime) 1.8.0の環境をセットアップする手順を纏めました。

2 JetPack 4.5

最新のJetPackは、下記でSDカードイメージ（jetson-nano-jp45-sd-card-image.zip 6GByte）としてダウンロードできます。

JetPack SDK

セットアップ後に、jtopで確認すると、JetPackのバージョンが確認可能です。

必須では有りませんが、DLRを使用してモデルを動作させる場合、少し大きなものだと、メモリ不足で止まってしまう事があるので。そこで、予め、スワップを拡張しておきます。

6Gのファイルを作成して、スワップを拡張している手順は以下のとおりです。

• wap領域（/swapfile）の確保

$ sudo dd if=/dev/zero of=/swapfile bs=1G count=6
6+0 records in
6+0 records out
6442450944 bytes (6.4 GB, 6.0 GiB) copied, 262.942 s, 24.5 MB/s

• すべてのswapプロセスを解除

＄sudo swapoff -a

• swapとして使うファイルの作成

$ sudo mkswap /swapfile
mkswap: /swapfile: insecure permissions 0644, 0600 suggested.
mkswap: /swapfile: warning: wiping old swap signature.
Setting up swapspace version 1, size = 6 GiB (6442446848 bytes)
no label, UUID=ef8bc090-901e-4561-b163-085c38981fb3

• swapファイルをアクティベート

$ sudo swapon /swapfile
swapon: /swapfile: insecure permissions 0644, 0600 suggested.

• 確認

$ grep SwapTotal /proc/meminfo
SwapTotal:       6291452 kB
$ free -h
              total        used        free      shared  buff/cache   available
Mem:           3.9G        958M        1.1G         22M        1.8G        2.7G           
Swap:          6.0G          0B        6.0G

3 DLR 1.8

SageMaker Neoでコンパイルされたモデルを使用するためには、DLR(Deep Learning Runtime)が必要です。
参考：https://github.com/neo-ai/neo-ai-dlr

JetPackのバージョンによっては、ビルド済みのモジュールがpipでインストール可能ですが、最新のJetPack4.5用のものはまだ無いため、今回は、コンパイルしました。
https://github.com/neo-ai/neo-ai-dlr/releases

コンパイルの手順は、以下のとおりです。

Building for NVIDIA GPU on Jetson Devices

DLRにはCMake3.13以降が必要となっていますが、予め入っていたのが、3.10.2だったので、ドキュメントに従い、CMakeをビルド・インストールしています。

$ cmake --version
cmake version 3.10.2

$ sudo apt-get install libssl-dev
 wget https://github.com/Kitware/CMake/releases/download/v3.17.2/cmake-3.17.2.tar.gz
$ tar xvf cmake-3.17.2.tar.gz
$ cd cmake-3.17.2
$ ./bootstrap
$ make -j4
$ sudo make install

$ cmake --version
cmake version 3.17.2

続いて、DLRのコンパイルとインストールです。

$ git clone --recursive https://github.com/neo-ai/neo-ai-dlr
$ cd neo-ai-dlr
$ mkdir build
$ cd build
$ cmake .. -DUSE_CUDA=ON -DUSE_CUDNN=ON -DUSE_TENSORRT=ON
$ make -j4
$ cd ../python
$ python3 setup.py install --user

importできることを確認しています。

$ python3
Python 3.6.9 (default, Oct  8 2020, 12:12:24)
[GCC 8.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import dlr

・・・パフォーマンスの向上のためのメトリック収集に関するアナウンス

>>> dlr.__version__
'1.8.0'

4 use_default_dlr=True

Amazon SageMaker NeoでJepson用にコンパイルしたモデルを解凍すると、以下のようなファイル構成になっています。

$ ls -la model
total 50260
drwxrwxr-x 2 nvidia nvidia     4096  2月 20 11:49 .
drwxrwxr-x 4 nvidia nvidia     4096  2月 20 11:49 ..
-rw-r--r-- 1 nvidia nvidia      682  2月  4 07:48 compiled.meta
-rw-r--r-- 1 nvidia nvidia      781  2月  4 07:48 compiled_model.json
-rw-r--r-- 1 nvidia nvidia       32  2月  4 07:48 compiled.params
-rwxr-xr-x 1 nvidia nvidia 44883920  2月  4 07:48 compiled.so
-rw-r--r-- 1 nvidia nvidia    14977  2月  4 07:48 dlr.h
-rw-r--r-- 1 nvidia nvidia  6539120  2月  4 07:48 libdlr.so
-rw-r--r-- 1 nvidia nvidia      647  2月  4 07:48 manifest

上のように展開したモデルを、dlr.DLRModel()で初期化すると、環境によって、エラーとなります。

model = dlr.DLRModel('model/', 'gpu')

Error message(s): libcudart.so.10.0: cannot open shared object file: No such file or directory

このエラーを回避するためには、モデルにバンドルされているlibdlr.soを無視するフラグが必要です。

model = dlr.DLRModel('model/', 'gpu', use_default_dlr=True)

参考：DLR Library (libdlr.so) could not be loaded #232

5 動作確認

下記のようなコードで、モデルが正常にロードでき、推論が可能になっているかどうかを確認することができます。

import numpy as np
import time
import os
import dlr
from dlr.counter.phone_home import PhoneHome
PhoneHome.disable_feature()

os.environ['TVM_TENSORRT_CACHE_DIR'] = '.'
# os.environ['TVM_TENSORRT_USE_FP16'] = '1'

model = dlr.DLRModel('model/', 'gpu', use_default_dlr=True)
img = np.random.rand(1, 3, 224, 224)

for _ in range(10):
    start = time.time() 
    model.run({'data': img})
    processing_time = time.time() - start
    print("{:.2f} sec".format(processing_time))

一度目の推論だけ、時間がかかっていますが、２回目以降は、安定した速度がでています。

$ python3 run.py
2021-02-20 11:51:41,623 INFO Forced to use default libdlr.so from /home/nvidia/.local/lib/python3.6/site-packages/dlr-1.8.0-py3.6.egg/dlr/libdlr.so
30.72 sec
0.02 sec
0.02 sec
0.02 sec
0.02 sec
0.02 sec
0.02 sec
0.02 sec
0.02 sec
0.02 sec

先のコードには、os.environ['TVM_TENSORRT_CACHE_DIR'] = '.'　が設定されていますが、この環境変数を設定することで、コンパイルされた内容がキャッシュされ、次回以降は、１回目の推論時間が短縮されます。

参考:Caching TensorRT Engines

ディレクトリを確認すると、キャッシュ（tensorrt_0_fp32.meta、tensorrt_0_fp32.plan）が確認できます。

$ ls -la
total 128792
drwxrwxr-x  2 nvidia nvidia     4096  2月 20 11:49 model
-rw-rw-r--  1 nvidia nvidia      476  2月 20 11:51 run.py
-rw-rw-r--  1 nvidia nvidia       88  2月 20 11:52 tensorrt_0_fp32.meta
-rw-rw-r--  1 nvidia nvidia 78326744  2月 20 11:52 tensorrt_0_fp32.plan

6 最後に

今回は、最新のJetPackとDLRを使用した推論環境のセットアップする手順を確認しました。

Amazon Sagemaker Neoで作成されたモデルが、Cudaのバージョン(10.0)に依存しているため、use_default_dlr=Trueを使用しない場合、JetPack4.2あたりが必要になってしまうのが、新しいJetPackを使用する場合の嵌りどころかも知れません。