概要

Amazon SageMaker Neo

とは

「どのフレームワークで学習したモデルでも、ボタン一つで多数の CPU / GPU / 専用チップに合わせた高速・小型バイナリを生成し、クラウドでもエッジでも同じモデルを再利用できる」仕組み
ref:https://docs.aws.amazon.com/ja_jp/sagemaker/latest/dg/neo.html

そもそも

複数のプラットフォームで推論を実行する機械学習モデルの最適化（ハードウェアの資源を効率的に利用して推論を高速化すること）は難しい。なぜなら、多くの最適化 / 推論加速ツールは特定の ML フレームワーク (PyTorch, TensorFlow など) に依存したり、特定のハードウェア (NVIDIA GPU, Intel CPU, Apple Neural Engine, Edge TPU 等) に特化しているため、環境が変わるたびに別々のツールチェーン・変換手順・パラメータ調整を人手で繰り返さざるを得ず、結果として開発速度や再現性が下がるため。もちろんAWSが開発したチップであっても、Inf1とInf2でも最適化のやり方は異なる。

これって...

Java VM (JVM)やコンテナ技術が目指すものと近い。異なるハードウェア/OSで同じコードを動かすというのはそもそも難しい。環境を吸収するような技術は時代が常に求めている。

なにをやっている？

フレームワーク固有グラフ（PyTorch, TF など）をフレームワーク非依存 IR（TVM Relay / TIR など）へ正規化する。内部的には Apache TVM 系技術を土台とするコンパイルパイプラインが行う代表的最適化はオペレータ融合 / 定数畳み込み / 静的メモリ計画 / データレイアウト変換など。

具体的には下記

オペレータ融合 (Operator Fusion): 複数の連続演算（例: 畳み込み→バッチ正規化→活性化）を一つのカーネルとしてまとめ、不要な中間メモリ書き出しとカーネル起動回数を減らす最適化。
定数畳み込み (Constant Folding): 実行時に変わらない計算（全入力が定数）をコンパイル時に先に計算し、結果を定数としてグラフに埋め込む最適化。
静的メモリ計画 (Static Memory Planning): 全中間テンソルの生存期間を解析し、同時に使われないメモリ領域を再利用することでピーク使用量と割り当て回数を減らす手法。
データレイアウト変換 (Layout Transformation): NCHW・NHWC・ブロッキング (NCHWc) など最適な次元並び／チャネル分割に再配置し、ハードウェア（CPU/GPU/アクセラレータ）のメモリアクセスやベクトル化を効率化する最適化。
これらを実施する理由：ディープラーニング推論は演算よりもしばしばメモリアクセス（DRAM⇄キャッシュ）やカーネル起動オーバーヘッドがボトルネックになるため、単純にフレームワークが逐次実行した場合は余計な中間バッファと起動回数で性能が低下する。これをまとめて減らし、かつデバイス特性（SIMD 幅やメモリ帯域）に沿わせる必要がある。

ref：https://aws.amazon.com/jp/blogs/news/amazon-sagemaker-neo-train-your-machine-learning-models-once-run-them-anywhere/

チュートリアル

やろうとしていること

軽量な画像認識モデルを一つ用意し、それをアウトプットし、SageMaker NEOを使ってモデルアーティファクトを生成し、ラズパイ５で動くことを確認する。

Amazon SageMaker Studio Code Editorにて画像認識モデルを用意する
Amazon SageMaker NEO を用いてモデルアーティファクトを生成する
ラズパイにモデルアーティファクトを設置し、画像認識を実施してみる

手順１：モデルの用意

どこで実行しても構いませんが、初めから実行環境が用意されているAmazon SageMaker Studio Code Editorが便利です。

セットアップはこちらを参照

準備ができたら、まずターミナルで下記を実行します。

pip install torch torchvision onnx

そして、次のpythonコードを実行します。
mobilenet_v2という学習済み画像認識モデルをロードして、特にカスタマイズせずエクスポートします。

import torch, torchvision

model = torchvision.models.mobilenet_v2(weights="IMAGENET1K_V1")
model.eval()

dummy = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model, dummy, "mobilenetv2.onnx",
    input_names=["input"],
    output_names=["logits"],
    opset_version=16,
    do_constant_folding=True,
    dynamic_axes=None  # 固定形状で Neo 向けに単純化
)

print("Exported mobilenetv2.onnx")

出力されたデータを圧縮します。

tar -czf mobilenetv2.tar.gz mobilenetv2.onnx

この圧縮されたデータをS3にアップロードしておきます。

手順2：モデルアーティファクトの作成

次に、エクスポートしたモデルをSageMaker NEOでモデルアーティファクトとして、ラズパイ５環境で利用できるようにコンパイルしていきます。

Amazon SageMaker AI へ移動
推論>コンパイルジョブ>コンパイルジョブの作成

下記のように設定します。

Job name: てきとう
Input location: 先ほどの圧縮ファイルを指定する
Framework: ONNX
Input configuration (DataInputConfig):

{"input":[1,3,224,224]}

対象デバイス:rasp4b
Output location: さっきと同じでOK
対象デバイス:rasp4bとしていますが、ラズパイ５で動くことを確認しています。
送信（8分くらい待つと出来上がる）

手順３：ラズパイで画像認識モデルの実行

モデルアーティファクトをラズパイで実行してみます。

コンパイルされたモデルとテスト用画像をラズパイに転送する
解凍する（tar -xzf xxxxx）
ラベルをダウンロードする

wget https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt -O imagenet_classes.txt

下記のスクリプトを実行する

import dlr, numpy as np
from PIL import Image
from dlr.counter.phone_home import PhoneHome

# 1. Disable phone home (optional)
PhoneHome.disable_feature()

# 2. Load model (current directory)
model = dlr.DLRModel(".", "cpu")

# 3. Preprocess
img = Image.open("test_img_rspi.png").convert("RGB").resize((224,224))
arr = (np.array(img).astype("float32") / 255.0).transpose(2,0,1)[None]

# 4. Warmup (optional)
for _ in range(3):
    model.run(arr)

# 5. Inference & Postprocess
outputs = model.run(arr)
logits = outputs[0]
pred = int(np.argmax(logits, axis=1)[0])
print("Predicted class index:", pred)

classes = [l.strip() for l in open("imagenet_classes.txt")]
print("label:", classes[88])  # 88番目（0始まり）

こちらの画像を入力してみました。（この画像はAmazon Nova Canvasで「極才色の鳥」とプロンプトを打って出てきたものです）

すると下記の結果となりました。

この鳥はmecawというそうですね（実際に調べたらそっくりでした）。

参考

Discussion