AWS Inf1でPyTorchモデルの推論を行う

AWSの推論環境の新たな選択肢として登場したInf1インスタンスを使い、PyTorchのモデルをコンパイル・推論させてみます。

概ねチュートリアル (https://github.com/aws/aws-neuron-sdk/blob/master/docs/pytorch-neuron/tutorial-compile-infer.md) をなぞった内容となります。

Inf1インスタンス

Inf1は昨年末のre:invent 2019で発表されたInferentiaチップ搭載のAWSインスタンスです。

• 2020/1/4現在、バージニア北部とオレゴンのみの提供となってます

aws.amazon.com

Inferentia

AWS Inferentiaは、AWSが開発した機械学習の推論に特化したチップです。サーバ用途を想定したもので、INT8で最大128 TOPS、FP16/BF16で最大64 TOPSを謳ってます。

サポートしている型などがCPUよりも制限されますので、モデルの変換やランタイムなどが必要となります。これらをフォローしているのがAWS NeuronというSDKです。
学習済みモデルをNeuron Executable File Format (NEFF) に変換し、Neuronラインタイム上で実行します。2019/1/4現在、PyTorch、TensorFlow、MXNetをサポートしてます。

github.com

チュートリアル

AWS Neuronのチュートリアルを参考に、PyTorchのモデル (ResNet50) のコンパイルと推論を行います。

• https://github.com/aws/aws-neuron-sdk/blob/master/docs/pytorch-neuron/tutorial-compile-infer.md

インスタンスの作成

オレゴンリージョンにコンパイル用 (inf1-compilation) と推論用 (inf1-inference) のインスタンスをそれぞれ作成していきます。

コンパイル環境 (inf1-compilation)

コンパイルで使うインスタンスの推奨スペックはc5.4xlarge以上とのことですので、c5.4xlargeインスタンスを作成します。

推論環境 (inf1-inference)

Deep Learning AMI (Ubuntu 18.04) を選択します。

ここではinf1.xlargeインスタンスを作成します。

sshで接続して、neuron-lsコマンドでInferentiaチップが利用可能かどうかチェックできます。

• inf1.xlargeインスタンスは1チップなので、1つだけ表示されています

$ /opt/aws/neuron/bin/neuron-ls
+--------------+---------+--------+-----------+-----------+------+------+
|   PCI BDF    | LOGICAL | NEURON |  MEMORY   |  MEMORY   | EAST | WEST |
|              |   ID    | CORES  | CHANNEL 0 | CHANNEL 1 |      |      |
+--------------+---------+--------+-----------+-----------+------+------+
| 0000:00:1f.0 |       0 |      4 | 4096 MB   | 4096 MB   |    0 |    0 |
+--------------+---------+--------+-----------+-----------+------+------+

コンパイル

inf1-compilationインスタンスにsshして作業します。

作業ディレクトリとしてinf1-tutorialを作って移動しておきます。

$ mkdir inf1-tutorial
$ cd inf1-tutorial/

OSパッケージインストール

依存するOSパッケージをインストールします。

$ sudo apt-get update
$ sudo apt-get install -y python3-venv g++

Python仮想環境の構築

venv環境をセットアップします。後ほどインストールするtorch-neuron用のリポジトリを追加しておきます。

$ python3 -m venv venv
$ source venv/bin/activate

(venv) $ pip install -U pip

(venv) $ vim /home/ubuntu/inf1-tutorial/venv/pip.conf
(venv) $ cat /home/ubuntu/inf1-tutorial/venv/pip.conf
[global]
extra-index-url = https://pip.repos.neuron.amazonaws.com

Pythonパッケージインストール

必要なPythonパッケージをインストールします。

気をつけないといけない点として、本家のtorchパッケージをインストールしてはいけません。torch-neuronに含まれており、競合してしまうからです。このため、torchvisionは--no-depsオプションを付けて本家torchをインストールしないようにしています。

(venv) $ pip install torch-neuron
(venv) $ pip install neuron-cc[tensorflow]
(venv) $ pip install pillow==6.2.2
(venv) $ pip install torchvision --no-deps
(venv) $ pip list | grep torch
torch-neuron      1.0.627.0
torch-neuron-base 1.3.0.1.0.34.0
torchvision       0.4.2

コンパイル

torchvisionからResNet50の学習済みモデルをロードし、コンパイルします。

コンパイルの実装自体は、torch.neuron.traceメソッドにモデルと入力サンプルを渡すだけです。TorchScriptと同じような感じですね。

import torch
import torch_neuron
from torchvision import models

# 学習済みモデルのロード
model = models.resnet50(pretrained=True)
model.eval()

# 入力サンプル
image = torch.zeros([1, 3, 224, 224], dtype=torch.float32)

# コンパイル
model_neuron = torch.neuron.trace(model, example_inputs=[image])

# ファイルに保存
model_neuron.save("resnet50_neuron.pt")

このPythonコードを実行します。Warningがバシバシ出てますが、Inf1インスタンス以外でコンパイルすると起こるもののようで、無視しても問題ないようです。

(venv) $ python trace_resnet50.py
INFO:Neuron:compiling module ResNet with neuron-cc
[E neuron_runtime.cpp:57] grpc server unix:/run/neuron.sock is unavailable. Is neuron-rtd running? Is socket /run/neuron.sock writable?
[E neuron_op_impl.cpp:52] Warning: Neuron runtime cannot be initialized; falling back to CPU execution
[E neuron_op_impl.cpp:53] Warning: Tensor output are ** NOT CALCULATED ** during CPU execution and only indicate tensor shape
...

(venv) $ ls -l resnet50_neuron.pt
-rw-rw-r-- 1 ubuntu ubuntu 41655304 Jan  4 01:38 resnet50_neuron.pt

出力されたファイルを、SCPなりS3なりで、inf1-inferenceインスタンスへ送ります。

推論

次にInf1インスタンス (inf1-inference) で環境構築し、inf1-compilationでコンパイルしたモデルを使って推論を行います。

$ mkdir inf1-tutorial
$ cd inf1-tutorial/

OSパッケージインストール -> inf1-compilationと同じ

Python仮想環境の構築 -> inf1-compilationと同じ

aws-neuron-runtimeとaws-neuron-toolsのインストール

Inferentiaを使って推論するためには、aws-neuron-runtimeとaws-neuron-toolsのインストールが必要です。

Deep Learning AMIではデフォルトでインストールされているようです。もしカスタムAMIなどを使っていてインストールが必要な場合は、以下のようにインストールします。

$ cat /etc/apt/sources.list.d/neuron.list
deb https://apt.repos.neuron.amazonaws.com bionic main
$ wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | sudo apt-key add -
$ sudo apt-get install aws-neuron-runtime
$ sudo apt-get install aws-neuron-tools

Pythonパッケージインストール

Pythonパッケージのインストールを行います。inf1-compilationと概ね同じなのですが、2点変更してます。

• neuron-cc[tensorflow]はコンパイル時のみ必要なので、ここでは行ってません
• チュートリアルに記載が無いのですが、torch-neuronはsixに依存してます
  • inf1-compilationではneuron-cc[tensorflow]にsixが含まれていました
  • リリースノートにもイシューとして記載されていますので、そのうち修正されるかと思います

(venv) $ pip install torch-neuron
(venv) $ pip install pillow==6.2.2
(venv) $ pip install torchvision --no-deps
(venv) $ pip install six

推論

推論コードは以下です。

import os
import time
import torch
import torch_neuron
import json
import numpy as np

from urllib import request
from torchvision import models, transforms, datasets

# 推論テスト用の画像をダウンロード
os.makedirs("./images", exist_ok=True)
request.urlretrieve(
    "https://raw.githubusercontent.com/awslabs/mxnet-model-server/master/docs/images/kitten_small.jpg",
    "./images/kitten_samll.jpg"
)

# ラベルをダウンロードして、IDとマッピング
request.urlretrieve(
    "https://s3.amazonaws.com/deep-learning-models/image-models/imagenet_class_index.json",
    "imagenet_class_index.json"
)
idx2label = []
with open("imagenet_class_index.json", "r") as read_file:
    class_idx = json.load(read_file)
    idx2label = [class_idx[str(k)][1] for k in range(len(class_idx))]

# Datasetを作成
eval_dataset = datasets.ImageFolder(
    os.path.dirname("./"),
    transforms.Compose([
        transforms.Resize([224, 224]),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        )
    ])
)

# 画像をロードしてバッチ化
image, _ = eval_dataset[0]
image = torch.tensor(image.numpy()[np.newaxis, ...])

# モデルのロード
model_neuron = torch.jit.load("./resnet50_neuron.pt")

# 推論
results = model_neuron(image)

# 推論結果トップ5を表示
top5_scores = results[0].sort(descending=True)[0][:5]
top5_idx = results[0].sort(descending=True)[1][:5]
top5_labels = [idx2label[idx] for idx in top5_idx]
print("Top 5 labels:", top5_labels)
print("Top 5 scores:", top5_scores)

このPythonコードを実行します。ラベルはCPU実行時と同等の結果が得られてますが、スコアの数値は少し異なります。コンパイルによって型の精度が float32 から float16 へ落ちているためです。

(venv) $ python infer_resnet50.py
Top 5 labels: ['tabby', 'Egyptian_cat', 'tiger_cat', 'lynx', 'tiger']
Top 5 scores: tensor([14.5000, 13.5000, 12.8125, 12.3750, 10.4375])

# CPU実行時の結果 (参考)
Top 5 labels: ['tabby', 'Egyptian_cat', 'tiger_cat', 'lynx', 'tiger']
Top 5 scores: tensor([14.5064, 13.5225, 12.8563, 12.4049, 10.5018], grad_fn=<SliceBackward>)

実際にInferentiaが使われているかどうかは、neuron-topコマンドで確認できます。このコマンドはランタイム上で実行中のモデルを表示します。

$ /opt/aws/neuron/bin/neuron-top
neuron-top - 2020-01-04 07:08:48
NN Models: 1 total, 1 running
Number of VNCs tracked: 1
0000:00:1f.0 Utilizations: Neuron core0 0.43%, Neuron core1 0.00%, Neuron core2 0.00%, Neuron core3 0.00%,
Model ID   Model Name                               UUID                               Node ID   Subgraph   Exec. Unit       Host Mem   Device Mem   Neuron core %
10006      1.0.5939.0+5849551057-/tmp/tmpfxyrlbuk   dc2079f02e9211ea89f19bebb3de27c8   1         0          0000:00:1f.0:0   10506624   146344448    0.43

速度比較

簡単にですが、Inferentia v.s. CPUで推論速度の比較を行います。上の推論コードのロード部分を以下のように書き換えました。

# ...省略...

if use_inferentia:
    # Inferentiaはコンパイル後のモデルを使う
    model_neuron = torch.jit.load("./resnet50_neuron.pt")
else:
    # CPUでは素のモデルを使う
    model_neuron = models.resnet50(pretrained=True)
    model_neuron.eval()

start_datetime = datetime.datetime.now()

# 同じ入力で1万回推論する
n = 10000
for i in range(1, n + 1):
    results = model_neuron(image)

    if i % 1000 == 0:
        print(f"{i}: {datetime.datetime.now()}")

finish_datetime = datetime.datetime.now()
print(finish_datetime - start_datetime)

結果は以下の表のとおりです。Inferentiaの方が15倍速い結果になりました。

CPUの推論速度がc5.largeの1コアと同等と仮定しますと、c5.largeの7〜8台分のスループットを1台でマークしていることになります。料金で比較しても3〜4割のコスト削減を期待できます。

アーキテクチャ	10000画像の推論時間 (sec)	1秒あたりのスループット (回)
Inferentia	45	222
CPU	667	15

まとめ

PyTorchモデルをInferentiaアーキテクチャ向けにコンパイルし、Inf1インスタンスで推論させてみました。CPU1コアの場合と比較して約15倍の高速なことを確認できました。

実際の利用にあたっては以下も勘案したほうが良さそうですが、推論環境の選択肢の1つとして検討するのは十分アリのようです。

• 複雑なモデルについてもさくっとコンパイルできるか
• TorchScriptでコンパイルして複数のコアで並列計算した場合のスループットと比較すべき
• インスタンスについてもP2/P3やG4ともコスト比較すべき

参考

• https://github.com/aws/aws-neuron-sdk/blob/master/docs/pytorch-neuron/tutorial-compile-infer.md
• https://github.com/awshlabs/reinvent19Inf1Lab
  • TensorFlowの場合はこちらのチュートリアルの方が良さそうです