GluonCVのモデルをSageMaker Neo + Jetson tx2 + DLRで動作させてみる

f:id:aptpod-tetsu:20200807224627j:plain 研究開発グループで機械学習に関係する仕事を担当している瀬戸です。前回は、fastaiで学習に使う関数をApache MXNetで真似してみた - aptpod Tech Blogを紹介させて頂きました。今回は、SageMaker Python SDKのMXNetで利用できるGluonCVのモデルを、SageMaker Neoでコンパイルし、Jetson tx2上でDLRを用いて動作させることができたので紹介したいと思います。

ツールの概要

Apache MXNet

Apache MXNet(以下、MXNet)は、TensorFlowPyTorchなど同じディープラーニングのフレームワークです。ここで紹介されているように、AWSのサービスであるAmazon SageMakerに利用されています。

GluonCV Toolkit

GluonCV Toolkit(以下、GluonCV)は、MXNetの持つパッケージの一つであるGluonをComputer Visionに特化する形で発展させたパッケージです。既知の有名な論文をベースにした学習済みモデルを提供するModel Zooや、新しいデータ拡張の関数などが提供されています。 Model Zooに設置されている学習済みモデルは、学習スクリプトその実行コマンドが公開されていて、再現性がある程度担保されるよう運用されています。使い方としては、Gluonベースの学習スクリプトに記述しているモデル部分やデータ拡張部分をGluonCVベースに書き換えるだけで使えて簡単に利用することが可能です。

Amazon SageMaker Neo & DLR

Amazon SageMaker NeoはAmazon SageMakerのサービスの一つでTVM StackDLRで動作するようにモデルをコンパイルするサービスです。TVMやDLRは、ディープラーニングフレームワークのモデルをGPUやCPUに最適化して実行してくれるツールです。

モデルを動かしてみる

1. GluonCVのModel Zooからモデルを保存する

Model Zooにあるモデルを扱うのが簡単なので、以下のスクリプトを使ってモデルをダウンロードしモデル保存します。モデルはtar.gzで圧縮して保存していますが、これはAmazon SageMaker Neoで扱うためです。この記事では、“ResNet152_v1d” を使って実験を行っていますが Model Zooのサイトの中から別のモデルを選択できます1

import tarfile
import argparse
import gluoncv
import mxnet as mx

from pathlib import Path


if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('--model-name', type=str, default='ResNet152_v1d')

    args = parser.parse_args()

    model_name = args.model_name

    # model_zooからモデルを取得
    net = gluoncv.model_zoo.get_model(model_name, pretrained=True)

    # モデルを保存するための前準備
    net.hybridize()
    # 画像の入力サイズが違う場合は変更する
    dummy = mx.nd.ones([1, 3, 224, 224])
    _ = net(dummy)

    # 保存先ディレクトリの設定と作成
    store_dir = Path(f'./{model_name}')
    if not store_dir.exists():
        store_dir.mkdir()

    # モデルの保存
    net.export(f'{store_dir}/{model_name}')

    # SageMakerNeoでコンパイルできるようにtar.gzで圧縮
    tar_name = './' + model_name + '.tar.gz'
    archive = tarfile.open(tar_name, mode='w:gz')
    archive.add(store_dir)
    archive.close()

2. Amazon SageMaker Neoでコンパイルする

任意のs3のバケットへモデルのtar.gzファイルを配置してAmazon SageMaker Neoでコンパイルをします。以下の画像のように、Amazon SageMakerサービスへログインしてコンパイルを行います。下図のようなGUI操作を行い、コンパイルされたモデルがtar.gzで吐き出されれば完了となります。

f:id:aptpod_tech-writer:20200803194415p:plain
Amazon SageMakerのコンパイル設定画面

3. Jetson TX2上のDLRを使って実行する

DLRをここを参考にインストールしたNVIDIA Jetson TX2上で以下のスクリプトを実行しました。前準備として、下記のスクリプトと同じディレクトリに、compiledフォルダを作成し、このフォルダ以下にコンパイル済みモデルを格納しておきます。

import sys
import numpy as np
from dlr import DLRModel
import datetime


if __name__ == '__main__':
    # モデルのロード
    device = 'gpu'
    model = DLRModel('compiled', device, 0)
    # 入力画像のサイズを指定
    image_size = 224

    times = []

    # 100回実行する
    for _ in range(100):
        # 入力データをランダムに生成する
        im = np.random.random([1, 3, image_size, image_size])
        b, h, w, c = np.array(im).shape

        input_data = {'data': im}

        # Predict 
        start_prediction_time = datetime.datetime.now()
        out = model.run(input_data)
       
        # 処理時間を保持 
        end_prediction_time = datetime.datetime.now()
        duration = end_prediction_time - start_prediction_time
        sys.stdout.write(f"\rprediction time duration: {duration}")
        sys.stdout.flush()
        times.append(duration.total_seconds())


    print()
    print('-------------------------------------------------------')
    print(f'total_time: {np.sum(times)}')
    print(f'mean_time: {np.mean(times)}')
    print(f'median_time: {np.median(times)}')
    print(f'std_time: {np.std(times)}')
    print('-------------------------------------------------------')

    print('finished prediction')

実行した結果は以下の画像のようになりました。推論速度は、今回の検証項目の内容ではありませんが試しに取得してみました。だいたい、1/0.053=18fpsぐらい出ているようです。

f:id:aptpod_tech-writer:20200803194512p:plain
ランダム入力による推論結果

また、FP16で推論した結果も載せておきます。おおよそ、中央値でみると1.7倍早くなっているようです。

f:id:aptpod_tech-writer:20200803194538p:plain
ランダム入力による推論結果(FP16)

まとめ

GluonCVのモデルを複数のツールを使ってJetson TX2上で動作させました。引き続き、機械学習に関するやってみたや業務内で検証した内容を紹介していこうと思います。


  1. 他のモデルでは試していないので動かない可能性があります。