MENU
AWS

AWS neuron Inf2での推論を試してみた

mickeyk

こんにちは、右も左も分からない未経験で入社し早くも9ヶ月が過ぎようとしているmickeyです。

今回は、AWS neuronについて深掘りして書いていこうと思います！

Contents
  1. AWS neuronとは
  2. 前準備
  3. コンパイル
  4. 推論精度
  5. 推論速度のベンチマーク
  6. コストパフォーマンス
  7. 結論

AWS neuronとは

AWS neuronとは、推論コストを下げつつパフォーマンスを最適化したいと考えるデベロッパーにとって魅力的なサービスです。
主な特徴は以下の通りです。

  • 最適化されたインフラ: AWS Inferentiaチップに最適化されており、高い推論性能と効率を実現します。
  • コスト効率: 高性能を保ちつつ、コストを抑えることができます。これにより大規模な機械学習モデルの推論を経済的に実行可能になります。
  • 広範なフレームワークサポート: TensorFlow, PyTorch, MXNetなどの主要な機械学習フレームワークに対応しており、既存のモデルを容易に移行できます。
  • エコシステム統合: Amazon EC2インスタンスとAmazon SageMakerなど、AWSの他のサービスとシームレスに統合されています。

今までトレーニング専用のインスタンスとしてTrn1、推論専用インスタンスとしてInf1が提供されていましたが、なんと今年の8月にinf2インスタンスが一部地域で使用可能となりました！
主に米国東部 (バージニア北部)、米国東部 (オハイオ)、米国西部 (オレゴン) リージョンで使用可能で、このインスタンスは従来のInf1よりも最小のコストで高いパフォーマンスを実現します。オンデマンドインスタンス、リザーブドインスタンス、スポットインスタンスとして、または Savings Plans の一部として購入可能です。（更なるInf2 インスタンスの詳細については、以下のAmazon EC2 Inf2 インスタンスのウェブページと AWS Neuron のドキュメントをご参照ください)。

昨今生成系AIが隆盛を極めておりますが、私も元々AIに興味があったということもあり、こちらのAWS neuron Inf2のリリースについても注目しておりました。なので今回Inf2インスタンスを実際に動かしてみた結果をまとめたいと思います。

前準備

幸いにも、AWSが公式にGitHub上でneuron用のサンプルjupyter notebookを提供しています
(GitHubのリポジトリは下記をご参照ください)。

GitHubには様々な推論モデルのJupyter Notebookが用意されており、中でも個人的に使用していたResNet50モデルを活用して、約11年前にAIブームを牽引した画像分類の実験での検証を行いたいと思います。このリポジトリには、LLM（Large Language Models）を使用するコードも含まれているため、興味のある方はぜひチェックしてみてください。使用するResNetモデルは、ImageNetで事前に訓練されたものです(ImageNetは、数百万枚の注釈付き画像を含む大規模なデータベースで、画像認識の研究や機械学習モデルの訓練に広く使われています)。この訓練済みモデルを使用して画像分類の予測を試みます。
まず、inf2のEC2インスタンスを立ち上げます。ここでの設定は以下の通りです。

  • インスタンスタイプ:
    • inf2.xlarge
  • リージョン: 米国東部 (バージニア北部)
    • リージョンコード: us-east-1
  • AMI: Deep Learning AMI Neuron PyTorch 1.13 (Ubuntu 20.04) 20231031
    • OS: Ubuntu
    • アーキテクチャ: x86_64

EC2インスタンスが立ち上がったらSSH接続し、jupyter notebookを立ち上げます(設定方法は割愛します)。
無事jupyter notebookが立ち上がったら、少し修正を加えつつ実際にコードを順次実行していき、色々検証していきます。

コンパイル

このコードでは、元のモデルをAWS Neuron SDKを利用して、Neuron対応のモデルにコンパイルしています。torch_neuronx.trace 関数を使い、モデルとサンプル入力画像を渡すことでコンパイルを行っています。コンパイルが完了した後、コンパイルにかかった時間を秒単位で出力し、その後、コンパイルされたモデルを TorchScript 形式で保存しています。保存されたファイル model.pt は、Neuronによる推論実行のためにデプロイされます。

import os
import urllib
from PIL import Image
import time

import torch
import torch_neuronx
from torchvision import models
from torchvision.transforms import functional

(途中省略)

# Create the model
model = models.resnet50(pretrained=True)
model.eval()
# Get an example input
image = get_image()

# Run inference on CPU
output_cpu = model(*image)

# Record the start time of the compilation.
start_time = time.time()
# Compile the model
model_neuron = torch_neuronx.trace(model, image)

# Record the end time of the compilation and calculate the elapsed time.
end_time = time.time()
compile_time = end_time - start_time
print(f"Compilation time: {compile_time} seconds")

# Save the TorchScript for inference deployment
filename = 'model.pt'
torch.jit.save(model_neuron, filename)

コンパイル時間を計測してみたところ、約24.4 [sec]でした。コンパイル時間は、使用するモデルの複雑さや入力のサイズによって大きく異なるため一概に断言できませんが、モデルが1度コンパイルされれば、その後は同じモデルを再コンパイルすることなく何度も推論を実行できるので、初期コストの時間と考えると個人的には短いと感じました。

推論精度

まず推論精度を見てみます。実際の画像分類に使用するサンプルとして以下のものを用意しました(balloonの画像です)。

コンパイルされたモデルを読み込んでNeuronを用いた推論を実行し、その結果をCPUでの推論結果と比較して、各々の上位5つの予測結果を表示しています。

import json

# Load the TorchScript compiled model
model_neuron = torch.jit.load(filename)

# Run inference using the Neuron model
output_neuron = model_neuron(*image)

# Compare the results
print(f"CPU tensor:    {output_cpu[0][0:10]}")
print(f"Neuron tensor: {output_neuron[0][0:10]}")

(途中省略)

# Lookup and print the top-5 labels
top5_cpu = output_cpu[0].sort()[1][-5:]
top5_neuron = output_neuron[0].sort()[1][-5:]
top5_labels_cpu = [id2label[idx] for idx in top5_cpu]
top5_labels_neuron = [id2label[idx] for idx in top5_neuron]
print(f"CPU top-5 labels:    {top5_labels_cpu}")
print(f"Neuron top-5 labels: {top5_labels_neuron}")

CPUとInf2インスタンスでの上位5つの結果は、以下のようになりました。

  • CPU top-5 labels: ['planetarium', 'water_tower', 'parachute', 'airship', 'balloon']
  • Neuron top-5 labels: ['planetarium', 'water_tower', 'parachute', 'airship', 'balloon']

最終的に表示された結果は、CPUとNeuronの推論で同じラベルが出力されていることから、Neuronを使った推論がCPUでの推論と同等であること、そして精度に問題がないことを確認することができます(top5にballoonが含まれています)。

推論速度のベンチマーク

推論の速度面について評価します。高いハードウェア利用率を達成し、システムのスループットを最大化するには複数のスレッドで並行して推論を行うことが不可欠です。NeuronCore アーキテクチャは、小バッチサイズでも効率良く動作し、GPUモデルと比べても優れた性能を発揮するよう最適化されています。今回は、最適化されたResNetモデルを2つのNeuronCores上で実行し、マルチスレッド推論のレイテンシーとスループットを計算するベンチマークテストを行います。

以下のコードを使用して、先ほどの画像に対して実際の推論速度を多角的に計測しました。

import time
import concurrent.futures
import numpy as np

def benchmark(filename, example, n_models=2, n_threads=2, batches_per_thread=1000):
    """
    Record performance statistics for a serialized model and its input example.

    Arguments:
        filename: The serialized torchscript model to load for benchmarking.
        example: An example model input.
        n_models: The number of models to load.
        n_threads: The number of simultaneous threads to execute inferences on.
        batches_per_thread: The number of example batches to run per thread.

    Returns:
        A dictionary of performance statistics.
    """

    # Load models
    models = [torch.jit.load(filename).to('cpu') for _ in range(n_models)]

    # Warmup
    for _ in range(8):
        for model in models:
            model(*example)

    latencies = []

    # Thread task
    def task(model):
        for _ in range(batches_per_thread):
            start = time.time()
            model(*example)
            finish = time.time()
            latencies.append((finish - start) * 1000)

    # Submit tasks
    begin = time.time()
    with concurrent.futures.ThreadPoolExecutor(max_workers=n_threads) as pool:
        for i in range(n_threads):
            pool.submit(task, models[i % len(models)])
    end = time.time()

    # Compute metrics
    boundaries = [50, 95, 99]
    percentiles = {}

    for boundary in boundaries:
        name = f'latency_p{boundary}'
        percentiles[name] = np.percentile(latencies, boundary)
    duration = end - begin
    batch_size = 0
    for tensor in example:
        if batch_size == 0:
            batch_size = tensor.shape[0]
    inferences = len(latencies) * batch_size
    throughput = inferences / duration

    # Metrics
    metrics = {
        'filename': str(filename),
        'batch_size': batch_size,
        'batches': len(latencies),
        'inferences': inferences,
        'threads': n_threads,
        'models': n_models,
        'duration': duration,
        'throughput': throughput,
        **percentiles,
    }

    display(metrics)

(一部割愛)

この計測により、下記の結果が得られました。
Filename: model.pt
Batch Size: 1
Batches: 2000
Inferences: 2000
Threads: 2
Models: 2
Duration: 0.989
Throughput: 2022.108
Latency P50: 0.983
Latency P95: 1.018
Latency P99: 1.032

少々分かりづらいですが、この結果から、2つのスレッドを使用して約0.989秒で2000回の推論を行い、平均レイテンシーは約0.983ミリ秒、スループットは2022.108推論/秒という高速なパフォーマンスを示しました。これは、モデルがリアルタイムアプリケーションでの使用に十分な速度を持っていることを意味し、使用環境によってはさらに最適化する余地があります。

コストパフォーマンス

推論の速度向上が認められる一方で、コストに関する疑問が浮かぶことと思います。
以下に、同等のvCPUとメモリリソースを持つ汎用CPUインスタンス、Neuron対応のインスタンス、及び基本的なGPUインスタンスの、米国バージニア北部リージョンでの時間当たりの料金を記載しました。

  • m5.2xlarge（汎用CPUインスタンス）: $0.384
  • inf2.xlarge（AWS Inferentia対応インスタンス）: $0.762
  • p2.xlarge（基本的なGPUインスタンス）: $0.90

汎用CPUインスタンスと比較して価格は高いものの、推論性能は大幅に向上します。基本的なGPUインスタンス（p2.xlarge）と比較して、inf2.xlargeは約16%安価です。このことは、AIモデルの推論に高性能な計算リソースを必要とする場合に、コスト効率の良い選択肢であることを意味します。

結論

今回の検証から得られる結果は、「Inf２インスタンスはコストパフォーマンスに優れている」ということです。リアルタイムアプリケーションに適する推論速度を提供しつつ、GPUインスタンスと比較してもコスト面で有利です。その為、AWSを利用して推論を行う業務において、Inf2インスタンスは十分に選択肢に入るべきだと考えられます。将来的には、さらに高性能なインスタンスが登場することが予想され、それらの検証も引き続き行っていく価値があると思われます。その際にはまた検証していきたいです！
これからもAIの情勢を見守りつつも、ITエンジニア1年目の終盤戦として基礎知識の体得も大事にしつつ頑張っていきます！

ここまで読んでいただいた皆さん、ありがとうございました。

AUTHOR
mickey
mickey
記事URLをコピーしました