Taste of Tech Topics

Acroquest Technology株式会社のエンジニアが書く技術ブログ

AWS SageMaker上でMobileNet SSDのモデルを学習する

皆さんこんにちは。@です
最近、AWSのSageMakerがいくつかアップデートがありました。個人的にはPyTorchのサポートがアツいですね。

さて、今回はSageMaker上で公式がサポートされていないアルゴリズムを学習する場合に、どのような方法があるのかを紹介していきます。
モデルはMobileNet SSDを題材として見ていきましょう。

SDK?コンテナ?

SageMaker上で自前のモデルを作成する方法は2つあります。

  1. SageMaker Python SDKを使う
  2. 自作のDockerコンテナを利用する

1の方法は、各フレームワークごとにSageMaker上で学習するためのスクリプト仕様が決まっていて、その仕様に沿って書けば学習・推論ができるようになっています。
例えばTensorflowであれば以下のページに仕様が載っています。
docs.aws.amazon.com

2の方法は、自分で学習・推論を行うDockerコンテナを定義して、SageMakerから利用する方法です。

今回は2のコンテナを利用する方法でモデルを作成していきます。(理由は後述します)
ただし、今回は学習してモデルを作成するところまでなので、推論は行いません。

今回作成するモデルについて

今回使用するMobileNet SSDは、物体検知のモデルであるSSDをより軽量にしたモデルです。
よくエッジデバイス上での物体検知に用いられます。アルゴリズムの詳細な内容の記載は省略します。
幸いコード自体はObject Detection APIのTensorFlow実装が公開されています。

github.com

今回は、Object Detection APIの学習スクリプトベースに開発します。
Object Detection APIに学習用スクリプトも用意されているため、可能な限り、再利用するために学習用のコンテナを利用します。
AWS SageMaker SDKで学習する場合は、インターフェースに併せて学習スクリプトを再実装する必要があります。
そのため、今回の場合ではObject Detection APIと比較して実装する量が多くなってしまいます。

SageMakerの学習時の仕組み

作業に入る前に、SageMakerで学習する際の仕組みを説明していきます。仕組みの概要は次の通りです。

  1. AWSでもともと用意されているか、または自分で定義したコンテナを立ち上げます。
  2. コンテナ起動時にSageMakerが、S3からコンテナ内に学習データを配置し、学習ジョブを開始します。
  3. 学習が終了後、出力されたモデルファイルはS3に配置され、コンテナを終了する。

f:id:acro-engineer:20180905010101p:plain

仕組みに沿って実装するために開発者が必要な作業は次の通りです。

  1. 学習を実行するイメージを作成する
  2. 作成したイメージをECR(Elastic Container Registry)にpushする
  3. SageMakerから学習ジョブを実行する

学習用コンテナを作成する

まずは学習用のイメージを作成していきます。
基本はObject Detection APIのコードを再利用しつつ、SageMakerのルールに合わせて変更していきます。
変更点は以下の3つです。

  1. コンテナ起動用のスクリプトを用意する
  2. 学習用データを/opt/ml/input/data以下に読みに行く
  3. モデルの出力先として/opt/ml/modelを指定する

コンテナ起動用のスクリプトを用意する

まず、1の変更点についてですが、SageMakerは学習用コンテナを、次のコマンドで起動します。

$docker <image名> train

SageMakerで学習するためにはコンテナ内部の環境変数PATHが読める箇所に、「train」の名前で学習スクリプトを配置する必要があります。PATHはDockerfileにて指定可能です。
そのため、今回は簡単に学習用のコードを実行するshellスクリプトを配置しました。

python object_detection/model_main.py \
    --alsologtostderr \
    --pipeline_config_path="ssdlite_mobilenet_v2_coco.config"  \
    --model_dir="/opt/ml/model"

より適切に実装するなら、Pythonで書いた学習用のコードを実装すべきなのですが、今回は簡単に試してみるということでこのような形にしました。
pipeline_config_pathには、MSCOCOデータセットでMobileNet SSDを学習するための設定ファイルを指定しています。
この設定ファイルはObject Detection APIに標準で用意されています。

学習用データを/opt/ml/input/data以下に読みに行く

次に2の学習用データの配置を設定します。
その前に抑えておく必要があるのは、SageMakerでは、コンテナとS3にあるデータはコンテナ上の/opt/mlディレクトリでやり取りされるということです。
ディレクトリ構造としては

/opt/ml/
    input/
       data/ ・・・S3からのデータが配置される
       config/ ・・・ハイパーパラメータなどの設定が配置される
    model/ ・・・ここに出力されたものがS3に出力される

のようになっています。
例えば、S3に配置されている学習用データは/opt/ml/input/data//に配置されます。
そのため、設定ファイルを以下のように変更することで、S3から配信される学習データを読むようにします。

※一部抜粋
train_input_reader: {
  tf_record_input_reader {
    input_path: "/opt/ml/input/data/train/mscoco_train.record"
  }
  label_map_path: "/opt/ml/input/data/train/mscoco_label_map.pbtxt"
}

ここではchannel名がtrainとなるように設定します。channel名については後ほど、SageMakerの学習ジョブの設定の際に説明します。
なので、後に出てくるSageMaker上でのジョブの設定時も学習データのchannel名をtrainに設定する必要があります。

モデルの出力先として/opt/ml/modelを指定する

また、3の出力するモデルのパスについても、train内でmodel_dir引数に/opt/ml/modelを指定することで、モデルファイルをS3に出力できます。

イメージを作成する

上記の3つを変更したうえで、Object Detection APIが動作するDockerイメージを作成します。
今回使用したDockerfileは以下のようになりました。

FROM python:2.7

RUN apt-get update
RUN apt-get install git -y
RUN pip install tensorflow

RUN curl -OL https://github.com/google/protobuf/releases/download/v3.3.0/protoc-3.3.0-linux-x86_64.zip;\
        apt-get install unzip;\
        unzip protoc-3.3.0-linux-x86_64.zip -d protoc3;\
        mv protoc3/bin/* /usr/local/bin/;\
        mv protoc3/include/* /usr/local/include/

RUN apt-get install python-pil python-lxml python-tk -y; \
            pip install Cython; \
            pip install jupyter; \
            pip install matplotlib; \
            pip install Cython; \
            pip install pillow; \
            pip install lxml; \
            pip install jupyter; \
            pip install matplotlib; 


ENV PYTHONUNBUFFERED=TRUE
ENV PYTHONDONTWRITEBYTECODE=TRUE
ENV PATH="/opt/program:${PATH}"

RUN mkdir /opt/program

RUN mkdir /opt/program/object_detection;\
        mkdir  /opt/program/slim
COPY train /opt/program/
COPY ssdlite_mobilenet_v2_coco.config /opt/program
COPY object_detection /opt/program/object_detection
COPY slim /opt/program/slim

WORKDIR /opt/program

RUN git clone https://github.com/cocodataset/cocoapi.git;\
        cd cocoapi/PythonAPI;\
        make;\
        cp -r pycocotools /opt/program/;

WORKDIR /opt/program

RUN /bin/bash -c "protoc object_detection/protos/*.proto --python_out=./"
ENV PYTHONPATH=$PYTHONPATH:/opt/program:/opt/program/slim

注意事項として、trainスクリプトを実行できるように環境変数PATHにtrainスクリプトが置いてあるパスを指定する必要があります。
この時のディレクトリ構成は以下のようになりました。Object Detection APIの実行に必要なライブラリと学習スクリプトであるtrainを配置しています。

 tf_container/
   object_detection/
   slim/
   train
   Dockerfile

ここで以下のコマンドを実行してDockerイメージを作成します。

$docker build -t object_detection:test .

以上でDockerイメージの作成は完了です。

作成したイメージをECRにpushする

次に作成したイメージをSageMaker上から指定できるようにECRにpushします。

#作成したイメージにタグ付けする
$docker tag object_detection:test <アカウントID>.dkr.ecr.<リージョン>amazonaws.com/object_detection:deploy3 
#作成したイメージをpushする
$docker push <アカウントID>.dkr.ecr.<リージョン>amazonaws.com/object_detection:deploy3

この際、認証周りのエラーが出た場合は以下を参考にECRの認証設定を行ってください。docs.aws.amazon.com

以下のようにリポジトリにタグ名が表示されていれば大丈夫です。
f:id:acro-engineer:20180702001100p:plain:w700

SageMakerから学習ジョブを実行する

それではいよいよ学習ジョブを実行していきます。
まずは、S3に学習用データを配置します。
f:id:acro-engineer:20180702001116p:plain:w700
注意事項としてbucket名にsagemakerという文字列が含まれている必要があります。

次に、学習ジョブを設定していきます。
SageMakerのトレーニングジョブの作成画面で以下のように設定していきます。
f:id:acro-engineer:20180702001451p:plain:w700

f:id:acro-engineer:20180702001944p:plain:w700

f:id:acro-engineer:20180702002510p:plain:w700

f:id:acro-engineer:20180702003003p:plain:w700

しばらくすると、ログに学習が進んでいる様子が出力されます。
f:id:acro-engineer:20180702004154p:plain:w800

学習が終了するとS3にモデルファイルが出力されます。
f:id:acro-engineer:20180702091601p:plain:w800

まとめ

今回はSageMaker上でコンテナからMobileNet SSDを学習してみました。
展望としては、epoch数などをハイパーパラメータから設定できればよりジョブを管理しやすくなりますね。
とはいえ、Dockerイメージを作成してしまえば、好きな環境で学習できるのでお手軽だと思いました。
それではまた。

Acroquest Technologyでは、キャリア採用を行っています。


  • ディープラーニング等を使った自然言語/画像/音声/動画解析の研究開発
  • Elasticsearch等を使ったデータ収集/分析/可視化
  • マイクロサービス、DevOps、最新のOSSを利用する開発プロジェクト
  • 書籍・雑誌等の執筆や、社内外での技術の発信・共有によるエンジニアとしての成長

 
少しでも上記に興味を持たれた方は、是非以下のページをご覧ください。

Kaggle Masterと働きたい尖ったエンジニアWanted! - Acroquest Technology株式会社のエンジニア中途・インターンシップ・契約・委託の求人 - Wantedlywww.wantedly.com