【加藤さん向け】オンプレで動かす機械学習パイプラインをSagemaker用に変更するときのポイント【社内共有】

せーの

2020.08.24

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

せーのでございます。

北海道はもうそろそろ秋の気配。みなさんも秋物の洋服をクリーニングに出そうとしたり、観葉植物の日当たりを工夫したりしてる時もともとオンプレやEC2で動かすために組んでいた機械学習のパイプラインをSagemaker用に書き直したいなと思うことって、よくありますよね。

でも、この作業にはSagemakerの勘所を押さえておく必要があります。今回はそんな書き換え作業時に押さえておくポイントを、がっつり社内向けに記述しておきます。

今回は特にターゲット層として機械学習関連を一緒に作業している「加藤さん」を念頭にこの記事を書いています。ですので「加藤さん」と同じくらいのバックグラウンドをお持ちの読者の方であればスッと入ってくるかと思います。

「加藤さん」像

私の考える「加藤さん」は

AWSのサービスについては基本押さえている
機械学習の基本的な用語や流れ（データセットなど）はわかる
Sagemakerを使ったことはない
dockerを使ったことはある
Sagemakerを使って機械学習をやってみたいが、とっかかりがわからない
扱うのはCV（Computer Vision）で、画像ファイルと、そのアノテーションファイル（json）をCOCO形式で使う

です。

皆様の中に「加藤さん」的要素のある方はこの先もお楽しみください。

オンプレとSagemakerの違い

それではまずはオンプレやEC2で機械学習環境を構築する場合と、Sagemakerを使う場合での違いについて書きます。

オンプレとSagemaker上での最大の違いは、オンプレは学習スクリプト、データセット、開発スクリプトを全て一つのインスタンスで行うのに比べて、Sagemakerでは学習をコンテナ上でスケールさせるため、それぞれのリソースを別のサービスで管理することです。

つまりオンプレではそれぞれのスクリプトが記述されているpythonコードをimportなどでつないで行けばOKですが

SagemekerではSagemaker上に配置されているjupyter notebookをコントロールにして、適切なECSコンテナを選び、そこにS3からデータセットを注入した上でECS上の学習スクリプトをスタートさせ、出来上がったモデルをS3にアウトプットさせる必要があります。

Sagemakerに予め用意されているTensorflowやCaffeなどのコンテナイメージを利用するときはAPI一発でECSがセットできますが、フレームワークのバージョンが古かったり新しすぎたりして見当たらない場合は、自分でカスタムコンテナを用意して、ECRに登録することでSagemakerから使えるようになります。

単純にコンテナを登録するだけならdockerをpushするだけなのですが、これをSagemaker SDKのAPIにそって動かすためには少しコツがあります。

(備忘録)docker imageをECRに登録する方法

ここで少し横道にそれますが、備忘録としてdocker imageをAmazon ECRに登録する方法をざくっと記述しておきます。

まずは普通にDockerFileを作成し、Dockerにbuildします。

docker build -t test:latest .

次にECRのレポジトリを作ります。

aws ecr create-repository --repository-name test --region ap-northeast-1

ecrのget-login-passwordを使ってdockerでログインします。

aws ecr get-login-password | docker login --username AWS --password-stdin <アカウントID>.dkr.ecr.ap-northeast-1.amazonaws.com/

buildしたimageを使ってECRにtagをつけます。

docker tag test <アカウントID>.dkr.ecr.ap-northeast-1.amazonaws.com/test:latest

ECRにimageをpushします。

 docker push <アカウントID>.dkr.ecr.ap-northeast-1.amazonaws.com/test:latest

ECRへのイメージ登録をおさらいしたところで、具体的なSagemakerでの書き方についてやっていきましょう。

オンプレ時からの変更点

S3からのインターフェースを作る

データセットなどのリソース類はS3からインプット/アウトプットします。具体的には、開発インスタンスにデータがある場合は

s3_path_train = sagemaker_session.upload_data(path="datasets/train/", bucket=bucket_name, key_prefix="train")

estimator.fit({'training': s3_path_train}, logs=True)

のような感じで開発インスタンスからS3にデータをアップロードして、そのパスを学習インスタンスに渡します。既にS3上にデータがある場合は

s3_path_train = sagemaker.session.s3_input('s3://{}/{}/train'.format(bucket_name, "train"), distribution='FullyReplicated',                                        content_type='image/jpeg', s3_data_type='S3Prefix')

estimator.fit({'training': s3_path_train}, logs=True)

のように場所を指定してあげるだけでOKです。

このパスを受ける側の学習スクリプトでは通常は

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--prepared-train-labels', type=str, required=True,
                        help='path to the file with prepared annotations')
    parser.add_argument('--train-images-folder', type=str, required=True, help='path to COCO train images folder')
...
...
...

    args = parser.parse_args()

のようにスクリプトを叩く時に引数として必要なデータのパスを付けてあげるのですが、Sagemaker SDKを使ってパスを渡しているので

from sagemaker_training import environment

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('--prepared-train-labels', type=str, required=True,
                        help='path to the file with prepared annotations')
...
...
...

    env = environment.environment()
    parser.add_argument('--train-images-folder', type=str, default=env.channel_input_dirs.get('training'), help='path to COCO train images folder')

...
...

    args = parser.parse_args()

のように[sagemaker_training(旧sagemaker_containers)]というライブラリを使ってデータをセットします。

ちなみに、sagemaker_containersを使って書くのであればこうなります。

import sagemaker_containers

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('--prepared-train-labels', type=str, required=True,
                        help='path to the file with prepared annotations')
...
...
...

    env = sagemaker_containers.training_env()
    parser.add_argument('--train-images-folder', type=str, default=env.channel_input_dirs.get('training'), help='path to COCO train images folder')

...
...

    args = parser.parse_args()

マネージドなコンテナイメージの場合

source_dirのrequirements.txtとスクリプト上でのインストールコマンドを使い分ける

SagemakerではEstimatorの定義時に[source_dir]という引数にディレクトリをセットすると、その下にある[requirements.txt]というファイルの中にあるライブラリをpipで自動でインストールしてくれます。

torch==0.4.1
opencv-python==3.4.0.14
numpy==1.16.0
torchvision==0.2.1
matplotlib==3.0.2
networkx==2.3
Pillow==6.2.2
Cython==0.28.2
pycocotools==2.0.0

しかし、requirements.txtは順番を保証してくれるものではありませんので、依存性のあるものを上から順番に記述してもインストールエラーになることがあります。例えばこの例で言えばpycocotoolsはCythonを使うのでCythonがインストールされていない状態でインストールするとエラーになります。そしてpycocotoolsの上にCythonがあるからと言って、Cythonのインストールを待ってpycocotoolsのインストールが開始するとは限りません。

この場合、ライブラリを使うスクリプトのinit時に入れてあげるときれいにいきます。今回で言えばCythonをrequirements.txtで入れてあげて、pycocotoolsはこの時点では入れません。

pycocotoolsはCOCO形式のデータセットを扱う時に必要になるので、そのライブラリが呼び出されたときにインストールします。

import argparse
import cv2
import os
from sagemaker_training import entry_point, environment, errors, runner

import torch
from torch.nn import DataParallel
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import transforms

from datasets.coco import CocoTrainDataset  ##&amp;amp;amp;lt;- ここ

import subprocess

pycocotools = subprocess.run(['pip', 'install', 'pycocotools'])

ちなみに、学習スクリプトを指定するentry_pointもこのsource_dirをルートにパスを指定します。つまり

estimator = PyTorch(entry_point='train.py',
                            source_dir='source_dir',
                            hyperparameters=hyper_param,
                            role=role,
                            framework_version='1.0.0',
                            train_instance_count=1,
                            train_instance_type='ml.p2.xlarge')

という定義をした場合、entry_pointに当たるtrain.pyは[source_dir/train.py]の位置にセットしておきます。

カスタムコンテナイメージの場合

次にカスタムでコンテナを作って学習させる場合のポイントです。カスタムでコンテナを作る場合は最初に何を入れても自由なので、requirements.txtにあるような環境設定やライブラリ群、学習スクリプトは予めコンテナに入れておき、動的に動かす必要がありそうなハイパーパラメータなどは外からいれるようにします。

基本のツリー構成

基本は

Amazon SageMakerで独自の学習/推論用コンテナイメージを作ってみる

を参考にしてみてください。噛み砕くと、Sagemaker SDKからデータやハイパーパラメータなどが渡されると、コンテナの中では自動的に[/opt/ml/XXXXXX]というフォルダが作られ、そこにセットされるので、学習ライブラリからはそのパスを使ってデータなどを指定してあげれば良い、ということになります。

まず最初にDockerFileをまるごと記述します。

FROM ubuntu:18.04
# timezone setting
ENV DEBIAN_FRONTEND=noninteractive
RUN apt-get update &amp;amp;amp;amp;&amp;amp;amp;amp; apt-get install -y tzdata
ENV TZ=Asia/Tokyo
ENV PATH="/opt/program:${PATH}"

COPY program /opt/ml/code
COPY init /opt/init

RUN chmod +x /opt/ml/code/train.py # trainファイルに実行権限を付与する

WORKDIR /opt/ml/code

RUN apt-get install libturbojpeg python3-tk python3-pip libsm6 libxrender1 libxext-dev -y
RUN pip3 install --upgrade pip
RUN pip3 install --upgrade setuptools
RUN pip3 install sagemaker-containers
RUN pip3 install -r /opt/init/requirements.txt
RUN pip3 install -r /opt/init/requirements2.txt

ENV SAGEMAKER_PROGRAM /opt/ml/code/train.py

コンテナにsagemakerライブラリを入れる

マネージドのコンテナイメージでは最初からsagemaker-trainingが入っていますが、カスタムでは自分で入れてあげる必要があるのでDockerFileに

RUN pip3 install sagemaker-training

のように忘れずにインストールしておきましょう。また、オンプレにてvirtualenvやanacondaを使って仮想環境を構築しているときは、ここでバラして入れてしまうとシンプルになります。

依存性のあるライブラリは

RUN pip3 install -r /opt/init/requirements.txt
RUN pip3 install -r /opt/init/requirements2.txt

のように複数回に分けて入れると順番に入ります。

トレーニング時の引数を変数でセットする

カスタムコンテナをSagemaker上から叩くにはいくつか方法があるのですが、DockerFileの[ENTRYPOINT]コマンドを使って

ENTRYPOINT ["python3", "train.py"]

のようにセットするとSagemaker SDKでセットしたhyperparametersが上手く引数にセットされないため、学習スクリプト側で大幅な変更が必要になります。 [SAGEMAKER_PROGRAM]という環境変数をセットしておくことで、SDKを使ったentry pointをサクッとセットすることができます。

ENV SAGEMAKER_PROGRAM /opt/ml/code/train.py

フレームワークのクラスではなく、汎用的なEstimatorクラスを使う

マネージドのコンテナイメージを使うときは

estimator = PyTorch(entry_point='train.py',
                            source_dir='source_dir',
                            hyperparameters=hyper_param,
                            role=role,
                            framework_version='1.0.0',
                            train_instance_count=1,
                            train_instance_type='ml.p2.xlarge')

のようにsagemaker.estimator.pytorchのような使用するコンテナのフレームワークと同じクラスを使いますが、カスタムコンテナの場合はそれを汎用的にしたestimatorクラスを使って書き直します。

estimator = Estimator(hyperparameters=hyper_param,
                        role=role,
                        image_name=image,
                        train_instance_count=1,
                        train_instance_type='ml.p2.xlarge')

こうすることでフレームワーク特化のAPIでは必須事項だったentry_pointやsource_dirのような引数をつけなくても良くなります。entry_pointの代わりに上に書いたようにSAGEMAKER_PROGRAMを使ってentry_pointを指定してあげます。

注意点

Sagemaker SDKとSDK2では引数の名前が違う

ここはハマりやすいです。Sagemaker SDKは2020年08月現在2.4.1が最新バージョンになっており、元ソースのgithubのmasterブランチや各ドキュメントもそのバージョンを元に書かれています。ですが、Sagemakerでjupyter notebookを開いてカーネルを選択するとバージョン1.7.1が入っています。

SDK ver.1とver.2では引数の名前（image_name => image_uri、など）が変わっていたり、必須項目が違っていたりするので、どちらのバージョンで構築するのかをまず決めて、それに合わせた組み方をする必要があります。 CFnなどで自動化しようとしている時は無理せずver.1を使ったほうが余計な手間はかからないかな、というのが個人的な感想です。