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）は、TensorFlowやPyTorchなど同じディープラーニングのフレームワークです。ここで紹介されているように、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 StackやDLRで動作するようにモデルをコンパイルするサービスです。TVMやDLRは、ディープラーニングフレームワークのモデルをGPUやCPUに最適化して実行してくれるツールです。

モデルを動かしてみる

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

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

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')