1 はじめに

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

Amazon SageMaker JumpStart（以下、JumpStart）は、TensorFlow Hub や PyTorch Hub に公開されているモデルをGUIで簡単にデプロイして利用できます。 以下は、PyTorch HubのResNet50でファイチューニングしてみた例です。

今回は、上記と同じ要領でResNet18から学習したモデルをSageMaker Neo（以下、Neo）でJetson Nano用に最適化して使用してみました。

前回、同じモデルをPyTorch上で使用した際の処理時間が、0.04sec 〜 0.07sec程度だったのに対し、今回は、0.02secとなっているので、倍以上の速度が出ていることになります。

2 コンパイル

JumpStartのファインチューニングで作成されたモデルは、PyTorchのモデルそのものですので、Neoでコンパイルが可能です。

(1) コンパイルジョブの作成

下記の設定で、コンパイルを開始します。

• アーティファクトの場所 : ファインチューニングで出力されたモデル
• データ入力設定 : {"data":[1,3,224,224]}
• 機械学習フレームワーク : PyTorch
• 対象デバイス : jetson_nano
• S3出力先 : モデルを出力するバケット

(2) No pth file found for PyTorch model.

しかし、先の要領でコンパイルすると、エラーとなってしまいます。

ClientError: InputConfiguration: No pth file found for PyTorch model. Please make sure the framework you select is correct.

ドキュメント上は、PyTorchのモデルは、「拡張子がptもしくはpthのものをtar.gzで圧縮したもの」となっているので、出力された圧縮ファイルがそのまま利用できるはずなのですが、何故か、拡張は、pthでないエラーとなってしまいました。

下記が、ファインチューニングで出力された圧縮ファイルです。

一旦、解凍し、拡張子をpthに変えて、tar.gzを作成し直しました。

% tree .
.
└── model
    └── model.pth

$ tar cfvz model.tar.gz model
a model
a model/model.pth

$ ls -la *.tar.gz
-rw-r--r--  1 hirauchi.shinichi  staff  41569044  2  4 06:09 model.tar.gz

(3) コンパイル完了

拡張子を変更すると、コンパイルは成功します。

3 JetPack 4.4

使用したJetPackは、4.4です。

DLRのコンパイル済みモジュールは、最新の4.5用がまだ無いため、一つ前の4.4をアーカイブから取得して利用しています。

https://developer.nvidia.com/embedded/jetpack-archive

https://developer.nvidia.com/jetpack-sdk-44-archive

4 DLR(Deep Learning Runtime)

DLRは、数日前に、バージョンが1.8となっていました。コンパイル済みのモジュールもJetPack4.4までは、公開されています。

https://github.com/neo-ai/neo-ai-dlr/releases

途中、numpyをsetupできるように、予めCythonを入れてから、dlrをインストールしています。
参考：Installing DLR

$ pip3 install Cython
Successfully installed Cython-0.29.21
$ pip3 install https://neo-ai-dlr-release.s3-us-west-2.amazonaws.com/v1.8.0/jetpack4.4/dlr-1.8.0-py3-none-any.whl

以下は、インストール状況を確認しています。

$ python3
Python 3.6.9 (default, Oct  8 2020, 12:12:24)
>>> import dlr
>>> dlr.__version__
'1.8.0'
>>> import numpy
>>> numpy.__version__
'1.19.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'] = '.'

model = dlr.DLRModel('model/', dev_type='gpu', use_default_dlr=True)
print("DLRModel() end")
print("model.get_input_dtypes(): {}".format(model.get_input_dtypes()))
print("model.get_input_names(): {}".format(model.get_input_names()))

img = np.random.rand(1, 3, 224, 224)
print(img.shape)
print("start")

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

5 コード

以下が、Jetson Nano上で、DLRを使用して推論しているコードです。

入力画像の変換や、出力値からSoftmax関数で信頼度を出力する処理は、ファインチューニンの学習時に使用されていたコードをそのまま利用させて頂いています。

# -*- coding: utf-8 -*-
import cv2
import numpy as np
import time
import os
import dlr
from dlr.counter.phone_home import PhoneHome
from torch.nn import functional as F
from torchvision import transforms
import torch

# Webカメラ
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)

# 入力画像の変換
RANDOM_RESIZED_CROP = 224
NORMALIZE_MEAN = [0.485, 0.456, 0.406]
NORMALIZE_STD = [0.229, 0.224, 0.225]
transform = transforms.Compose(
	[
		transforms.ToPILImage(), # Numpy -> PIL
		transforms.Resize(RANDOM_RESIZED_CROP),
		transforms.CenterCrop(RANDOM_RESIZED_CROP),
		transforms.ToTensor(),
		transforms.Normalize(NORMALIZE_MEAN, NORMALIZE_STD),
	]
)

class_names = ["ASPARAGUS","CRATZ","NOIR","OREO","PRETZEL"]

def main():
	cap = cv2.VideoCapture(GST_STR, cv2.CAP_GSTREAMER)

	# フォーマット・解像度・FPSの取得
	width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
	height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
	print("width:{}　height:{}".format(width, height))

	# SHARE = 224
	PhoneHome.disable_feature()
	os.environ['TVM_TENSORRT_CACHE_DIR'] = '.'
	model = dlr.DLRModel('model/', 'gpu', use_default_dlr=True)
	print("DLRModel() end")
	print("model.get_input_dtypes(): {}".format(model.get_input_dtypes()))
	print("model.get_input_names(): {}".format(model.get_input_names()))

	while True:

		# カメラ画像取得
		_, frame = cap.read()
		if(frame is None):
			continue
		image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
		image = transform(image) # 入力形式への変換
		image = image.unsqueeze(0) # バッチ次元の追加 (3,224,224) => (1,3,224,244)

		start = time.time()
		outputs = model.run({'data': image})
		processing_time = time.time() - start
		
		batch_probs = F.softmax(torch.from_numpy(outputs[0]), dim=1)
		batch_probs, batch_indices = batch_probs.sort(dim=1, descending=True)
		
		for probs, indices in zip(batch_probs, batch_indices):
			name = class_names[indices[0].int()]
			confidence = probs[0]
			str = "{} {:.1f}% {:.2f} sec".format(name, confidence, processing_time)
			print(str)
			cv2.putText(frame, str, (10, HEIGHT - 100), cv2.FONT_HERSHEY_PLAIN, 4, (255, 255, 255), 5, cv2.LINE_AA)

		cv2.imshow('frame', frame)

		if cv2.waitKey(1) & 0xFF == ord('q'):
			break

	# VideoCaptureオブジェクト破棄
	cap.release()
	cv2.destroyAllWindows()

if __name__ == '__main__':
	main()

6 最後に

今回は、JumpStartのファインチューニングで作成されたPytorchモデルをNeoで最適化してJetson Nanoで使用してみました。 コードなどは、Pytorchのものと殆ど変わりなく利用でき、倍以上の速度がでてしまうNeoは、強力だと思います。