[Amazon SageMaker] イメージ分類のモデルをNeoで最適化して、Jetson Nano+OpenCV+Webカメラで使用してみました

2020.06.29

1 はじめに

CX事業本部の平内(SIN)です。

前回、Amazon SageMaker(以下、SageMaker)の物体検出(組み込みアルゴリズム)を、SageMaker Neo(以下、Neo)で最適化して、Jetson Nanoで利用してみました。

今回は、イメージ分類(組み込みアルゴリズム)について、確認してみました。

最初に、動作を確認している様子です。GPUがフルに回っていますが、約0.1秒で推論できています。

2 モデル

使用したモデルは、下記で作成したものです。

17種類の商品を回転台に乗せて動画撮影したデータから、イメージ分類のモデルが作成されています。

3 SageMaker Neo

下記の諸元で、上記のモデルを最適化しています。

  • ジョブ名: ic-SYOHIN17-jetson-Nano-001(任意です)
  • データ入力値: {"data": [1, 3, 224, 224]}
  • 機械学習フレームワーク: MXNet
  • 対象デバイス: jetson_nano

ステータスが、COMPLATEになったら完了です。コンパイルは数秒で完了します。

出力モデルの名前は、model-jetson_nano.tar.gzとなります。

ダウンロードして展開してみると、圧縮ファイルの内容は、以下の通りでした。(注:最近、Neoの出力するファイル構成が変わったと思います。)

model-jetson_nano
├── compiled.meta
├── compiled.params
├── compiled.so
├── compiled_model.json
└── model-shapes.json

4 Jetson Nano

Neoで最適化したモデルを使用するには、DLR(Deep Learning Runtime)のセットアップ必要です。

手順については、下記で紹介させて頂いた要領と同じです。
参考:[Amazon SageMaker] 動画から生成したデータセットで商品棚の商品を検出してみました

$ python3
Python 3.6.9 (default, Apr 18 2020, 01:56:04)
[GCC 8.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import dlr
>>> dlr.__version__
'1.2.0'
>>>

[2020/07/08追記]なお、DLRセットアップ中、setup.pyで、「No module named 'setuptools'」となりましたので、pipでインストールしました。

$ python3 setup.py install --user
Traceback (most recent call last):
  File "setup.py", line 3, in <module>
    from setuptools import setup, find_packages
ModuleNotFoundError: No module named 'setuptools'
$ sudo apt-get update
$ sudo apt-get install python3-pip
$ pip3 --version
pip 9.0.1 from /usr/lib/python3/dist-packages (python 3.6)
$ pip3 install setuptools

また、モデルのサイズにもよりますが、DLRで最初に推論する際に、下記のログが出力され、TensorRTエンジンのビルドが行われます。 メモリが不足する場合、このログの後、ハングアップしてしまいます。

[08:15:29] /home/nvidia/work2/neo-ai-dlr/3rdparty/tvm/src/runtime/contrib/tensorrt/tensorrt_module.cc:80: Building new TensorRT engine for subgraph tensorrt_0

今回のモデルを実行するために、SWAPをデフォルト値の2Gから6Gに増やしています。

$ grep SwapTotal /proc/meminfo
SwapTotal:       6291452 kB

5 コード

確認に使用したコードは、以下のとおりです。

先に、Neoで作成したモデル(model-jetson_nano.tar.gz)は、./modelに展開されています。

import cv2
import numpy as np
import dlr
import time

MODEL_PATH = './model'
CLASSES = ['PORIPPY(GREEN)', 'OREO', 'CUNTRY_MAM', 'PORIPPY(RED)', 'BANANA'
           , 'CHEDDER_CHEESE', 'PRETZEL(YELLOW)', 'FURUGURA(BROWN)', 'NOIR'
           , 'PRIME', 'CRATZ(RED)', 'CRATZ(GREEN)', 'PRETZEL(BLACK)', 'CRATZ(ORANGE)'
           , 'ASPARA', 'FURUGURA(RED)', 'PRETZEL(GREEN)']
SHAPE = 224

DEVICE_ID = 0
WIDTH = 800
HEIGHT = 600
GST_STR = ('v4l2src device=/dev/video{} ! video/x-raw, width=(int){}, height=(int){} ! videoconvert ! appsink').format(DEVICE_ID, WIDTH, HEIGHT)

def putText(frame, y, text):
    x = 20
    size = 1.5
    width = 3
    cv2.putText(frame, text, (x, y), cv2.FONT_HERSHEY_SIMPLEX, size, (255,255,255), width, cv2.LINE_AA)

def main():
    # Version
    print("Ver dlr:{} cv2:{}".format(dlr.__version__, cv2.__version__))

    # Video Initialize
    cap = cv2.VideoCapture(GST_STR, cv2.CAP_GSTREAMER)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    print("fps:{} width:{} height:{}".format(fps, width, height))

    # Model Initialize
    model = dlr.DLRModel(MODEL_PATH, 'gpu')
    print("input_dtypes: {}".format(model.get_input_dtypes()))
    print("input_names: {}".format(model.get_input_names()))

    print("model OK")
    while(True):

        _, frame = cap.read()
        if(frame is None):
            continue
        frame = frame[0 : int(height), 0 : int(height)] # 横長の長方形 => 正方形

        # 入力画像生成
        img = cv2.resize(frame, dsize=(SHAPE, SHAPE)) # height * height => 224 * 224
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # BGR => RGB
        img = img.transpose((2, 0, 1)) # 224,244,3 => 3,224,224
        img = img[np.newaxis, :] # 3,224,224 => 1,3,224,224
        print("img.shape: {}".format(img.shape))

        # 推論
        start = time.time() # 時間計測
        out = model.run({'data': img})
        processing_time = time.time() - start
        print(processing_time)

        # 表示
        prob = np.max(out)
        index = np.argmax(out[0])
        putText(frame, height-80, CLASSES[index])
        putText(frame, height-30, "{:.3f} {:.3f}sec".format(prob, processing_time))

        # 画像表示
        cv2.imshow('frame', frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
main()

6 最後に

SageMaker組み込みアルゴリズムの物体検出をNeoでコンパイルするためには、MXNetへの変換が必要でしたが、イメージ分類の場合は必要ありません。

Jetson NanoへのDLRのインストールとメモリ容量に問題がなければ、特に詰まる所は無いかも知れません。

エッジ側で、0.1秒で推論できれば、それなりに利用範囲はあると思いました。