見出し画像

【Part8】学習モデルをハードウェアに最適化したコンパイル(Amazon SageMaker Neo)

AWSでエッジコンピューティング環境を作る Part8です。

学習モデルの軽量化Distillerを、Part7で書きました。今日は、そのモデルをエッジ端末(ラズパイ)上で動かす準備を行います。Amazon SageMakerNeoを用いると、対象のハードウェアに合わせて最適化されたコンパイルをよしなに行ってくれるので、とても良いです。

・PyTorchのモデル保存とONNX形式
・Amazon SageMaker Neoでのコンパイル
・エラーの格闘記録

あたりをお話します。

PyTorchでの学習

前回のDistillerの流れもあり、PyTorchを引き続き使っていきます。PyTorchを使ってみたいという方は、以前私がチュートリアルをやってみた内容を記事に上げているので、ご覧になってください。

REFチュートリアルやってみた

また、こちらの記事では、学習済みモデルを読み込んで推論&再学習する方法が丁寧に書かれているので、本格的に機械学習を進めていきたい方は必読です。

REFPyTochでpre-trainedモデルを再学習する

手順(ざっくり)

1. 前回作ったdistillerのモデルを読み込み、ONNX形式に変換する
2. ONNX形式に変換されたモデルをS3に投げ込む
3. Amazon SageMaker Neoを用いてコンパイルする

これだけ!簡単そうですよね。ハマったは最後に書いて行くとして、手順、書いていきましょう。

1. モデル読み込み・変換

前回作ったafter_compression.pth.tarを用います。後述しますが、sagemaker neoがPyTorchの最新バージョンに対応していない(にもかかわらず、distillerは最新バージョンを要求してくる...)ので、DLモデルの共通フォーマットであるONNX形式に変換します。変換コードが↓。compress_testがあるフォルダ内に配置しています。

import os
from distiller.models import create_model
import torch
import tarfile

modelpath = os.path.join("compress_test","after_compress.pth.tar")
onnxfile = "model.onnx"

### pyTorchモデルの読み込み
model_pth = torch.load(modelpath)
### モデルの重みとか、パラメータが"state_dict"に格納されている
parameters =model_pth["state_dict"]

### distillerで同じ箱を作ってやる
model = create_model(False, "cifar10", "simplenet_cifar")
### さっき読み込んだパラメータをはめ込む
model.load_state_dict(parameters)

### ONNXモデルとして出力
# verbose=Trueとしておくと、詳細がわかる
dummpy_input = torch.randn(1,3,32,32, device='cpu')
torch.onnx.export(model, dummpy_input, onnxfile,verbose=True)

# 出力したONNXファイルをtar.gzに圧縮
with tarfile.open('model.tar.gz', 'w:gz') as f:
   f.add(onnx)

実行すると、verbose=Trueにしているので、モデルの詳細を見ることができます(なぜこうしているかは後述)。以下のようなログが出てくれば成功です。

graph(%0 : Float(1, 3, 32, 32),
     %conv1.weight : Float(6, 3, 5, 5),
     %conv1.bias : Float(6),
     %conv2.weight : Float(16, 6, 5, 5),
     %conv2.bias : Float(16),
     %fc1.weight : Float(120, 400),
     %fc1.bias : Float(120),
     %fc2.weight : Float(84, 120),
     %fc2.bias : Float(84),
     %fc3.weight : Float(10, 84),
     %fc3.bias : Float(10)):
 %11 : Float(1, 6, 28, 28) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[5, 5], pads=[0, 0, 0, 0], strides=[1, 1]](%0, %conv1.weight, %conv1.bias), scope: Simplenet/Conv2d[conv1]
 %12 : Float(1, 6, 28, 28) = onnx::Relu(%11), scope: Simplenet/ReLU[relu_conv1]
 %13 : Float(1, 6, 14, 14) = onnx::MaxPool[kernel_shape=[2, 2], pads=[0, 0, 0, 0], strides=[2, 2]](%12), scope: Simplenet/MaxPool2d[pool1]
 %14 : Float(1, 16, 10, 10) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[5, 5], pads=[0, 0, 0, 0], strides=[1, 1]](%13, %conv2.weight, %conv2.bias), scope: Simplenet/Conv2d[conv2]
 %15 : Float(1, 16, 10, 10) = onnx::Relu(%14), scope: Simplenet/ReLU[relu_conv2]
 %16 : Float(1, 16, 5, 5) = onnx::MaxPool[kernel_shape=[2, 2], pads=[0, 0, 0, 0], strides=[2, 2]](%15), scope: Simplenet/MaxPool2d[pool2]
 %17 : Tensor = onnx::Constant[value=  -1  400 [ Variable[CPUType]{2} ]](), scope: Simplenet
 %18 : Float(1, 400) = onnx::Reshape(%16, %17), scope: Simplenet
 %19 : Float(1, 120) = onnx::Gemm[alpha=1, beta=1, transB=1](%18, %fc1.weight, %fc1.bias), scope: Simplenet/Linear[fc1]
 %20 : Float(1, 120) = onnx::Relu(%19), scope: Simplenet/ReLU[relu_fc1]
 %21 : Float(1, 84) = onnx::Gemm[alpha=1, beta=1, transB=1](%20, %fc2.weight, %fc2.bias), scope: Simplenet/ReLU[relu_fc1]
 %22 : Float(1, 84) = onnx::Relu(%21), scope: Simplenet/ReLU[relu_fc2]
 %23 : Float(1, 10) = onnx::Gemm[alpha=1, beta=1, transB=1](%22, %fc3.weight, %fc3.bias), scope: Simplenet/ReLU[relu_fc2]
 return (%23)

2. S3に保存

ONNXに変換したモデルができたので、S3に保存しましょう。こちらはバケットを適当に作り、モデルをアップロードするだけです。S3の作り方はPart3あたりを参考にしてください。

画像1

3. SageMakerNeoでコンパイル

S3に配置したモデルをラズパイで使えるようにコンパイルしましょう。AWSのSageMakerNeoを用います。SageMaker⇨コンパイルジョブを選択し、作成を初めましょう。

画像2

ジョブ名はアカウント&リージョンで一意に定まるように、適当に決めてください。また、IAMロールでは、先程作成したS3にアクセスできるようにロールを作成しましょう。

画像3

つぎは入力設定です。まず、S3に格納したモデルの場所を設定します。次がデータ入力設定。そしてフレームワークをONNXとして完了です。

画像4

ここで、データ入力設定が重要です。データの型が間違っていると永遠にエラーを吐きます。ここでさっきの変換時に使ったverboseが活きてきます。出力結果の一部を再掲します。

graph(%0 : Float(1, 3, 32, 32),
     %conv1.weight : Float(6, 3, 5, 5),
     %conv1.bias : Float(6),
     %conv2.weight : Float(16, 6, 5, 5),
     %conv2.bias : Float(16),
     %fc1.weight : Float(120, 400),
     %fc1.bias : Float(120),
     %fc2.weight : Float(84, 120),
     %fc2.bias : Float(84),
     %fc3.weight : Float(10, 84),
     %fc3.bias : Float(10)):

この一番はじめの%0 : Float(1, 3, 32, 32),が入力の型です。これを参考に、{"0" : [1, 3, 32, 32]}と入力します。間違えないようにしましょう。

入力が終われば、最後は出力の設定です。出力場所はさっきのS3と同じにして、compiled/以下に保存されるようにしましょう。対象となるハードウェアはラズパイなので、それ用に最適化してもらいます。すべて入力し終われば、コンパイルを開始しましょう。

画像5

成功すれば下図のようになります。

画像6

また、S3にはコンパイル後のモデルが入っています。
これをGreengrassを用いてデプロイすれば、学習モデルをラズパイで使うことができます!

画像7

Amazon SageMaker Neoで学習モデルのコンパイル格闘記録

画像8

実は、全然スムーズにコンパイルできたわけではありません。PyTorchのモデルをコンパイルしようとすると、実に様々なエラーがでました。格闘記録をエラーメッセージとともに書いておきます。

☑ tar.gzで保存する必要がある。

ClientError: InputConfiguration: Unable to untar input model. Please confirm the model is a tar.gz file

☑ input shapeが決まっているようだ

ClientError: InputConfiguration: Invalid PyTorch model or input-shape mismatch. Make sure that inputs are lexically ordered and of the correct dimensionality.

Neoコンパイルエラー時のトラブルシューティング

ここにシェイプタイプの例が書かれているので参考に。気をつけるのは、そのフレームワークがNHWC形式なのかNCHW形式なのかを把握しておくこと

☑ NHWC : Tensorflow
☑ NCHW : MXNet, PyTorch, ONNX
N : Number
H : Height
W : Width
C : Channel

でないと、エラーが出る

☑ PyTorchの互換性が無い

ClientError: InputConfiguration: Invalid PyTorch model or input-shape mismatch. Make sure that inputs are lexically ordered and of the correct dimensionality.

PyTorchだと永遠にエラーがでる。
調べると、PyTorch1.1.0以降には未対応(2019/9/1現在)の模様。使いたい場合は0.4.0にデグレードするか、ONNXに変換するかという選択肢になる。ONNXとは、ニューラルネットの共通フォーマットのひとつで、様々なライブラリを共通のモデルの型に変換できる。

ONNXOpen Neural Network Exchange
ニューラルネットの共通フォーマットの一つ。
☑ ONNX モデルへの変換 (エクスポート)
・Caffe2
・PyTorch
・CNTK
・Chainer
☑ ONNX モデルを用いた推論 (インポート)
・Caffe2
・CNTK
・MXNet
・TensorFlow
・Apple CoreML
・TensorRT (ただしサンプルコードが未公開)

REF1: ONNX公式
REF2: Qiita ONNX形式のモデルを扱う

ONNXにエクスポートすることでとりあえず解決

☑ ONNXにおいて、input shapeのはじめのkeyが違う

Neoコンパイルエラー時のトラブルシューティング
上記リンクの例の通り入力をおいて、ONNXにしても、まだエラーがでた。

ClientError: InputConfiguration: Incompatible ONNX mode

どうやら入力に使っているkeyが違っているらしい。手順3で説明したように、変換時にverbose=Trueにすると、わかる。これで調べて、input shapeをきちんと書き直したところ、見事Finish

長かった...

まとめ

ここまでで、学習モデルをラズパイ向けにコンパイルし、Greengrassにデプロイする準備ができました。途中PyTorchに慣れたり、互換性等々の問題もあって少し苦労しましたが、無事にできてよかったです。

いよいよ次回、Greengrassでデプロイしてエッジコンピューティングを実行していきましょう。コンパイルした推論モデルを動かすDLR(Deep Learning Runtime)を導入し、エッジ推論ができるか試していきます。ではではっ

Part 7. モデル圧縮専用ライブラリDistiller
Part 9. GreengrassとDLR

この記事が気に入ったら、サポートをしてみませんか?気軽にクリエイターを支援できます。

note.user.nickname || note.user.urlname

サポートいただけると励みになります! よろしくおねがいします!!

ありがとうございます!!がんばります
11
i Smart Technologies 執行役員CTO『現場目線×デジタル手法で製造業を再構築。』面白いことを求め、 IoT/cloud/データ分析/機械学習/MASなど技術全般やってます。東大航空⇒NTT研究所⇒iSTC

コメント2件

日々、濃い記事を読ませていただいております!ありがとうございます。
なぜこんなにも熱心に投稿されるのですか?
いつもありがとうございます!
半年くらいたった自分が同じことで困らないようにです。そして人目に晒されると思うとちゃんと書こうとするので!笑
コメントを投稿するには、 ログイン または 会員登録 をする必要があります。