AWSの推論用インスタンスInf1について（part1）

NTTPCではAIを用いたサービスを複数提供しており、私もそのうちの1つである姿勢推定による動作解析APIプラットフォームサービス 『AnyMotion 』というサービスの開発に携わっています。
私のチームではインフラ基盤としてAWSを用いており、今回は推論用インスタンスの1つであるInf1インスタンスに焦点を当てて検証を行いました。

検証内容をInf1インスタンスの導入を検討されている方を対象として記事にします。記事は2部構成で、1部ではInf1の概要と使い方について記載します。2部では速度検証とコスト試算について記載します。
本記事は1部であり、目次は次の通りです。

1. Inf1インスタンス概要
2. コンパイル
3. 推論実行
4. AWSとの連携とまとめ

まずは、Inf1インスタンスについて説明します。
Inf1インスタンスは、推論用のインスタンスで、AWS独自開発の高性能な機械学習推論用プロセッサであるInferentiaを搭載しています。
InferentiaチップにはNeuronコアが4個搭載されており、その上でモデルが動作します。
Inf1インスタンスの概要図は次の通りです。^*1

インスタンスサイズによって、以下のように1枚から16枚のInferentiaチップが搭載されています。^*2

Neuronコア上でモデルを動作させるには、モデルをNeuron SDKというツールで適した形にコンパイルする必要があります。
TensorFlow、PyTorch、MXNetなどの一般的な機械学習のモデルからコンパイルが可能です。また、コンパイル後のモデルの実行や、プロファイリングにも同SDKを用います。
次のような利用フローとなります。

推論実行時に、Neuronコアの使用方法はカスタマイズ可能です。
例えばinf1.xlargeの場合、4つのNeuronのうちの1つのみを用いてモデルを動作させることや、4つのNeuronコアそれぞれで1つのモデルを動作させること、4つのNeuronコアにまたがる大きなモデルを動作させること、4つのNeuronコアのうち2つを1つのモデルで利用し、残りを別のモデルで利用することなどが可能です。

また、6xlarge以上のInf1インスタンスでは、複数のInferentiaチップにまたがってNeuronコア数を指定することも可能です。

（1）コンパイルするモデルについて

AnyMotionにて実際に使用している姿勢推定モデル（AnyMotionモデル）を用いて、コンパイルと推論を行いました。

AnyMotionモデルは人間の骨格位置を推定してその座標値を返すというTensorflowモデルです。
入出力を抜粋したものを次に示します。

inputs['images'] tensor_info:
    dtype: DT_FLOAT
    shape: (-1, 256, 192, 3)
outputs['coords']

（2）検証環境の作成

当検証ではコンパイルから実行までを同一のインスタンス上で行います。
Neuron SDKがプリインストールされているDeep Leaning AMIをAMIに指定してinf1.xlargeインスタンスを作成します。環境構築の詳しい手順はAWS公式のハンズオン資料の環境セットアップの章をご確認下さい。

当記事執筆時点（2021年6月28日）での情報を記載します。ソフトウェアやパッケージは頻繁に更新されるため、ご利用の際は、最新の環境をご利用下さい。

EC2 Instance:：inf1.xlarge
Deep Learning AMI ：Deep Learning AMI (Ubuntu 18.04) Version 46.0
Neuron SDK：version 1.4.1.0+737cbb69a

Neuronのバージョンは次のコマンドで確認しました。

$ pip list | grep neuron
neuron-cc                          1.4.1.0+737cbb69a
tensorboard-plugin-neuron          2.1.0.0
tensorflow-neuron                  1.15.5.1.4.0.0

（3）コンパイルの実施

次のコードでコンパイルを実施しました。

import shutil
import tensorflow as tf
import tensorflow.neuron as tfn

model_dir = '<コンパイル前のsaved_model配置しているディレクトリのパス>'
compiled_model_dir = '<コンパイル後のsaved_modelを配置するディレクトリのパス>'

tfn.saved_model.compile(model_dir, compiled_model_dir, batch_size = 1, dynamic_batch_size = True)

コンパイル時のパラメータについて次に記載します。

実行結果は次の通り、コンパイルが完了しました。

INFO:tensorflow:Successfully converted /home/ubuntu/model/1 to compiled_mode

（4）問題の発生と解決

コンパイルは完了しましたが、作成したモデルで推論を実行したところ処理に失敗し、次のログが出力されました。

tensorflow.python.framework.errors_impl.InvalidArgumentError: 2 root error(s) found.
            (0) Invalid argument: no batch-dimension found on input tensor neuron_op_e47de1c2b8587fc20/_3:0 with shape []
               [[{{node coords/neuron_op_d87260133335f92}}]]
               [[coords/neuron_op_d87260133335f92/_15]]
            (1) Invalid argument: no batch-dimension found on input tensor neuron_op_e47de1c2b8587fc20/_3:0 with shape []
               [[{{node coords/neuron_op_d87260133335f92}}]]
          0 successful operations.

調査の結果、グラフの一部（coordsというノード）がうまくコンパイル出来ていないことが分かりました。
no_fuseという、指定した部分のコンパイルを行わないオプションを付けてコンパイルを行うことで上記問題を解決しました。

import shutil
import tensorflow as tf
import tensorflow.neuron as tfn

model_dir = "<コンパイル前のsaved_modelを配置しているディレクトリのパス>"
compiled_model_dir = '<コンパイル後のモデルを配置するディレクトリのパス>'

with tf.Session(graph=tf.Graph()) as sess:
    tf.saved_model.loader.load(sess, ['serve'], model_dir)
    no_fuse_ops = [op.name for op in sess.graph.get_operations() if 'coords' in op.name]
tfn.saved_model.compile(model_dir, compiled_model_dir, batch_size=1, 
dynamic_batch_size=True, no_fuse_ops=no_fuse_ops)

検証中、コンパイル関連で詰まった部分2箇所を次のNoteに記載します。

Note1. 対応していないノードがある場合

記事執筆時点とはコンパイラのバージョンが違うため異なる挙動でしたが、検証時点（2020年12月）でも（4）と同様の問題が発生しました。環境は次の通りで、（3）のコードでコンパイルを行いました。

EC2 Instance:：c5.4xlarge
Deep Learning AMI ：Deep Learning AMI (Ubuntu 16.04) Version 38.0
Neuron SDK：version 1.1.7.0+6d41df113

その結果コンパイル時に次のエラーが発生しました。調査の結果、（4）と同じ原因でコンパイルに失敗していました。

ValueError: slice index 1 of dimension 0 out of bounds. for 'coords/strided_slice_2' (op: 'StridedSlice') with input shapes: [0], [1], [1], [1] and with computed input tensors: input[1] = <1>, input[2] = <2>, input[3] = <1>.

no_fuseオプションを付けてコンパイルを行うことで問題を解決しました。

import shutil
import tensorflow as tf
import tensorflow.neuron as tfn

model_dir = "<コンパイル前のsaved_modelを配置しているディレクトリのパス>"
compiled_model_dir = '<コンパイル後のモデルを配置するディレクトリのパス>'

with tf.Session(graph=tf.Graph()) as sess:
    tf.saved_model.loader.load(sess, ['serve'], model_dir)
    no_fuse_ops = [op.name for op in sess.graph.get_operations() if 'coords' in op.name]
tfn.saved_model.compile(model_dir, compiled_model_dir, batch_size=1, dynamic_batch_size=True, no_fuse_ops=no_fuse_ops)

Inf1は新しい技術であり、アップデートが早いのでログの内容やエラーを出力する状況など挙動が変わることもあります。検証の際にはどのバージョンを利用しているかを記録して行うことをお勧めします。

Note2. 入力サイズが可変なモデルの場合の設定

当記事のモデルは入力サイズが固定のモデルですが、入力サイズが可変のモデルについても（3）と同様のコードで行ったところ、コンパイルに失敗しました。

Neuronでは入力が可変なモデルは利用できないため、コンパイルする際にサイズを決定する必要があります。
model_feed_dictオプションにて入力サイズを指定することで問題を解決できました。

import shutil
import tensorflow as tf
import tensorflow.neuron as tfn

model_dir = "<このパイル前のsaved_modelを配置しているディレクトリのパス>"
compiled_model_dir = '<コンパイル後のモデルを配置するディレクトリのパス>'

# インプットの形式を定める
batch_size = 1
model_feed_dict = {'inputs': np.zeros([batch_size, 224, 224, 3])}

tfn.saved_model.compile(model_dir, compiled_model_dir, batch_size=1, dynamic_batch_size=True, model_feed_dict=model_feed_dict)

（1）実行

コンパイルバッチサイズ1のモデルを推論インスタンス上に配置し、推論バッチサイズ2で処理回数1000回分の画像の推論にかかる時間を計測しました。

NEURONCORE_GROUP_SIZESで使用するNeuronの数を指定できます。次のコードではNEURONCORE_GROUP_SIZESに4x1を指定することでNeuronコアを4つ使用しています。

import os
import shutil
import numpy as np
import tensorflow as tf
import time

# 設定
os.environ['NEURONCORE_GROUP_SIZES'] = '4x1'
batch_size = 2
num_inferences = 1000

# モデルの読み込み
compiled_model_dir = '<コンパイル後のモデルを配置しているディレクトリのパス>'
predictor_inferentia = tf.contrib.predictor.from_saved_model(compiled_model_dir)

# warm up
model_feed_dict={'images': np.zeros([batch_size, 256, 192, 3])}
infa_rslts = predictor_inferentia(model_feed_dict)
print(infa_rslts["coords"])

# 推論実行
start = time.time()
for _ in range(num_inferences):
    infa_rslts = predictor_inferentia(model_feed_dict)
elapsed_time = time.time() - start

# 処理時間とスループットの表示
print('処理画像枚数:{:>6}[images], 処理時間:{:6.2f}[sec], スループット:{:8.2f}[images/sec]'.format(num_inferences*batch_size, 
elapsed_time, num_inferences*batch_size / elapsed_time))

結果は次の通りです。

処理画像枚数:  2000[images], 処理時間:  7.10[sec], スループット:  281.54[images/sec]

推論中のNeuronコアの使用状況はneuron-topというコマンドで確認出来ます。
確認したところ、Neuronコアが4つ使われていることが確認出来ました。
ただ、Neuronコアの使用率が30%台のためまだ余裕がありそうです。

（2）改善

推論バッチサイズを2から8に変えて検証を行いました。
次に示す通りNeuronコアの使用率は80%付近になり、スループットも向上していることが確認出来ました。

処理画像枚数:  8000[images], 処理時間:  19.33[sec], スループット:  413.82[images/sec]

Neuron SDKでのコンパイル時のエラーを解決することに苦労したので、第1部ではコンパイル時に困ったこととその解決策、コンパイル時のオプションについて詳しく記載しました。

処理速度についてはAWSの他のインスタンスと比べて向上しそうだと感じました。
第2部では、推論バッチサイズ、コンパイルバッチサイズなどのパラメータを変化させて、Inf1以外のインスタンスとも比較しながら速度検証とコスト試算を行います。

宜しければそちらもご一読下さい。

第2部：『AWSの推論用インスタンスInf1について（part2）』

また1部2部通して、問題解決及び性能検証に当たってはAWSに支援を頂戴しました。特にInf1及びNeuron SDKは新しい技術ということもあり、AWSからの支援は大変助かりました。この場を借りて改めてお礼申し上げます。

^*1Inf1 ハンズオン資料：https://ec2-inf1.workshop.aws/ja/

^*2Inf1 公式ページ：https://aws.amazon.com/jp/ec2/instance-types/inf1/