Amazon SageMakerの組み込みアルゴリズムのセマンティックセグメンテーションを試してみる – Amazon SageMaker Advent Calendar 2018

こんにちは、データインテグレーション部の大澤です。

今年も残すところ後一ヶ月となりました！
というわけで今日からデータインテグレーション部機械学習チームで「クラスメソッド Amazon SageMaker Advent Calendar」をお送りします！
Amazon SageMaker Examplesをやってみるのを中心にSageMakerでできることをご紹介していければと思います。よろしくお願いします。

1日目は、先日新しくSageMakerの組み込みアルゴリズムとして追加された、セマンティックセグメンテーション(Semantic Segmentation)についてお伝えします！

セマンティックセグメンテーションとは

セマンティックセグメンテーションは画像に写っているものを推論する問題の１つです。似たものとして画像分類と物体検出があります。それぞれの内容と違いはざっくり次のような感じです。

• 画像分類: 画像内のもののラベルを推論します。どういうラベルのものが画像に含まれているのかが分かります。
• 物体検出: 画像内のもののラベルと検出したものを囲むバウンディングボックス(長方形)の位置を推論します。どういうラベルのものが画像に含まれているのかが、バウンディングボックス単位で分かります。
• セマンティックセグメンテーション: 画像内のもののラベルと検出したもののピクセル単位での位置を推論します。どういうラベルのものが画像に含まれているのかが、ピクセル単位で分かります。

今回は扱うのはこの中のセマンティックセグメンテーションです。教師データとしては、画像の各ピクセルがどのラベルかという情報が含まれた行列データになります。 例えば、このような画像であれば....

教師データは次のようになります。

やってみる

では実際にサンプルノートブックに沿ってやってみたいと思います!

• amazon-sagemaker-examples/semantic_segmentation_pascalvoc.ipynb at master · awslabs/amazon-sagemaker-examples

セットアップ

モデルの学習やエンドポイントの作成時に使用するIAMロールを定義します。 今回はノートブックインスタンスに関連づけられているIAMロールを使用します。

%%time
import sagemaker
from sagemaker import get_execution_role

role = get_execution_role()
print(role)

sess = sagemaker.Session()

学習用データや学習時の出力結果のモデルアーティファクトを配置するS3のバケット名とオブジェクトキーのプリフィックスを定義しておきます。

# リージョンごとのデフォルトのバケットを使用する。必要に応じて変更する
bucket = sess.default_bucket()

prefix = 'semantic-segmentation-demo'
print(bucket)

学習やエンドポイントに使用するコンテナイメージの名前を取得します。

from sagemaker.amazon.amazon_estimator import get_image_uri
training_image = get_image_uri(sess.boto_region_name, 'semantic-segmentation', repo_version="latest")
print (training_image)

今回はこのイメージを使用します。

501404015308.dkr.ecr.ap-northeast-1.amazonaws.com/semantic-segmentation:latest

データセットのダウンロード

今回使用するPASCAL VOC2012のデータセットをダウンロードします。PASCAL VOCはセマンティックセグメンテーションなどのコンピュータービジョンの大会です。 教師用画像が1464枚、バリデーション用画像が1449枚あり、ラベルは21種類あります。

• PASCAL VOC2012

%%time

# データセットをダウンロード
!wget -P /tmp http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar

# 展開する
!tar -xf /tmp/VOCtrainval_11-May-2012.tar && rm /tmp/VOCtrainval_11-May-2012.tar

容量が大きいため、ダウンロード時間も結構かかります。

データの準備

教師画像とバリデーション用画像とそれらのアノテーション画像を適切な入力形式に対応させます。 今回の場合の入力データのざっくりとした条件は以下の通りです。

• 各入力データはデータチャネル名(trainとかvalidationなど)のディレクトリ内に配置する。
• 教師画像とバリデーション用画像はjpg形式で、それらのアノテーション画像は8bitの単一チャネルのpng形式である必要があります。
• アノテーション画像の名前は対応する教師画像(もしくはバリデーション用画像)と合わせる必要があります。
• 必要に応じてラベルマップを作成して入力できます。ラベルマップは人間が読める文字列とアノテーションで各ピクセルで設定されているラベル番号の対応関係です。

今回の場合は基本的に条件を満たしているので、S3で配置する形に合わせてデータチャネルごとにディレクトリを分ければOKです。以下のような配置にします。

s3://bucket_name/prefix/
    |
    |- train
                 |
                 | - 0000.jpg
                 | - coffee.jpg
    |- validation
                 |
                 | - 00a0.jpg
                 | - bananna.jpg
    |- train_annotation
                 |
                 | - 0000.png
                 | - coffee.png
    |- validation_annotation
                 |
                 | - 00a0.png
                 | - bananna.png
    |- label_map
                 | - train_label_map.json
                 | - validation_label_map.json

上のような配置になるように各画像をコピーします。

import os
import shutil

# S3にデータを出力する時の配置を模してディレクトリを作成する
VOC2012 = 'VOCdevkit/VOC2012'
os.makedirs('train', exist_ok=True)
os.makedirs('validation', exist_ok=True)
os.makedirs('train_annotation', exist_ok=True)
os.makedirs('validation_annotation', exist_ok=True)


# 教師画像の一覧を作成する
filename = VOC2012+'/ImageSets/Segmentation/train.txt'
with open(filename) as f:
    train_list = f.read().splitlines()

# バリデーション用画像の一覧を作成
filename = VOC2012+'/ImageSets/Segmentation/val.txt'
with open(filename) as f:
    val_list = f.read().splitlines()

# 教師画像(jpg)とそのアノテーション画像(png)を作成したディレクトリへコピーする
for i in train_list:
    shutil.copy2(VOC2012+'/JPEGImages/'+i+'.jpg', 'train/')
    shutil.copy2(VOC2012+'/SegmentationClass/'+i+'.png','train_annotation/' )

# バリデーション用画像(jpg)とそのアノテーション画像(png)を作成したディレクトリへコピーする
for i in val_list:
    shutil.copy2(VOC2012+'/JPEGImages/'+i+'.jpg', 'validation/')
    shutil.copy2(VOC2012+'/SegmentationClass/'+i+'.png','validation_annotation/' )

今回ラベルマップは使用しませんが、ノートブックにはラベルマップの作成処理も書かれています。利用する際には参考にすると良さげです。

import json
label_map = { "scale": 1 }
with open('train_label_map.json', 'w') as lm_fname:
    json.dump(label_map, lm_fname)

各データチャネルのパスを定義しておきます。

train_channel = prefix + '/train'
validation_channel = prefix + '/validation'
train_annotation_channel = prefix + '/train_annotation'
validation_annotation_channel = prefix + '/validation_annotation'
# label_map_channel = prefix + '/label_map'

Amazon SageMaker Ground Truthのラベリングジョブによって生成されたマニフェストファイルも入力として使えるようです。 そのほかの事項や詳細についてはセマンティックセグメンテーションのドキュメントをご覧ください。

S3へのデータのアップロード

各データチャネルのデータをS3へそれぞれアップロードします。今回はラベルマップは不要ですが、必要に応じてコメントアウトを外してアップロードを実行してください。

%%time
sess.upload_data(path='train', bucket=bucket, key_prefix=train_channel)
sess.upload_data(path='validation', bucket=bucket, key_prefix=validation_channel)
sess.upload_data(path='train_annotation', bucket=bucket, key_prefix=train_annotation_channel)
sess.upload_data(path='validation_annotation', bucket=bucket, key_prefix=validation_annotation_channel)
# sess.upload_data(path='train_label_map.json', bucket=bucket, key_prefix=label_map_channel)

モデルアーティファクトの出力場所も定義しておきます。

s3_output_location = 's3://{}/{}/output'.format(bucket, prefix)
print(s3_output_location)

モデルの学習

セマンティックセグメンテーションのモデルの学習を行います。

まずは学習を行うためのエスティメーターを設定します。学習を行うインスタンスタイプやインスタンス数、最大学習時間などの学習の基礎情報を設定します。

ss_model = sagemaker.estimator.Estimator(training_image,
                                         role, 
                                         train_instance_count = 1, 
                                         train_instance_type = 'ml.p3.2xlarge',
                                         train_volume_size = 50,
                                         train_max_run = 360000,
                                         output_path = s3_output_location,
                                         base_job_name = 'ss-notebook-demo',
                                         sagemaker_session = sess)

次にハイパーパラメータを設定します。 ハイパーパラメータの詳細についてはドキュメントをご覧ください。

ss_model.set_hyperparameters(backbone='resnet-50', # エンコーダ
                             algorithm='fcn', # デコーダー。そのほかの設定として、'psp'や'deeplab'があります。
                             use_pretrained_model='True', # 事前学習させたモデルを使用
                             crop_size=240, # ランダムに切り取る画像のサイズ
                             num_classes=21, # クラス数
                             epochs=10, # １つのデータを何回学習させるか。ノートブックによると、30エポック位がPASCAL VOCのデータセットに合っているようですが、今回はデモのため10で留めます。
                             learning_rate=0.0001,   # 学習率
                             optimizer='rmsprop', # 最適化関数。そのほかには'adam', 'rmsprop', 'nag', 'adagrad'があります。
                             lr_scheduler='poly', # 学習率を変化させる方法、そのほかの選択肢として'cosine'や 'step'があります。
                             mini_batch_size=16, # ミニバッチ（一度に学習させる単位）のデータ数
                             validation_mini_batch_size=16, # バリデーション時のミニバッチサイズ
                             early_stopping=True, # 早期の学習停止を有効にするかどうか.
                             early_stopping_patience=2, # mIoUが増えなくなってから、何エポック後に学習を終了するか
                             early_stopping_min_epochs=10, # 最低何エポック学習させるか(指定したエポック数までは早期学習は無効)
                             num_training_samples=num_training_samples) # 学習用データの数

データチャネルの定義を作成します。

# 各データチャネルのフルパスを定義
s3_train_data = 's3://{}/{}'.format(bucket, train_channel)
s3_validation_data = 's3://{}/{}'.format(bucket, validation_channel)
s3_train_annotation = 's3://{}/{}'.format(bucket, train_annotation_channel)
s3_validation_annotation = 's3://{}/{}'.format(bucket, validation_annotation_channel)

distribution = 'FullyReplicated'

# s3_inputデータを作成
train_data = sagemaker.session.s3_input(s3_train_data, distribution=distribution, 
                                        content_type='image/jpeg', s3_data_type='S3Prefix')
validation_data = sagemaker.session.s3_input(s3_validation_data, distribution=distribution, 
                                        content_type='image/jpeg', s3_data_type='S3Prefix')
train_annotation = sagemaker.session.s3_input(s3_train_annotation, distribution=distribution, 
                                        content_type='image/png', s3_data_type='S3Prefix')
validation_annotation = sagemaker.session.s3_input(s3_validation_annotation, distribution=distribution, 
                                        content_type='image/png', s3_data_type='S3Prefix')
# 各データチャネルの定義
data_channels = {'train': train_data, 
                 'validation': validation_data,
                 'train_annotation': train_annotation, 
                 'validation_annotation':validation_annotation}

学習を実行します。学習には20分程度かかります。東京リージョンでp3.2xlargeの場合だと3ドル弱かかるので注意が必要です。

ss_model.fit(inputs=data_channels, logs=True)

マネジメントコンソールのジョブ詳細ページから精度やロスなどの変化を見ることができます。 以下の画像はバリデーションデータにおける各エポックでのピクセル毎の精度です。

エンドポイントの作成

学習させたモデルをホストするエンドポイントを作成します。エンドポイントが立ち上がるまで数分かかります。

ss_predictor = ss_model.deploy(initial_instance_count=1, instance_type='ml.c4.xlarge')

推論

エンドポイントが立ち上がったら、そのエンドポイントに画像を投げて推論してみます。

まずはテスト用の画像をPexelsからダウンロードします。

!wget -O test.jpg https://images.pexels.com/photos/242829/pexels-photo-242829.jpeg
filename = 'test.jpg'

推論用に画像サイズを変更し、バイト配列に変換します。

import matplotlib.pyplot as plt

import PIL

# 推論用に画像サイズを変更します
im = PIL.Image.open(filename)
im.thumbnail([800,600],PIL.Image.ANTIALIAS)
im.save(filename, "JPEG")


%matplotlib inline
plt.imshow(im)
plt.axis('off')
with open(filename, 'rb') as image:
    img = image.read()
    img = bytearray(img)

推論に使う画像はこの画像です。

推論時に指定が必要なパラメータを設定し、推論を実行します。

• content_type: 推論用エンドポイントに投げるデータの形式です。値としては'image/jpeg'です。学習用データと同じ形式のjpeg画像です。
• accept: 推論結果のデータ形式です。'image/png'と'application/x-protobuf'が選択可能です。
  • image/png: 教師データと同じ形式のpng画像
  • application/x-protobuf: プロトコルバッファ形式。ピクセルごとに各クラスの確率(自信度)が返されます。

まずはacceptタイプをpng画像の場合から試します。

%%time 
ss_predictor.content_type = 'image/jpeg'
ss_predictor.accept = 'image/png'
return_img = ss_predictor.predict(img)

推論にかかる時間はおよそ300m秒程度でした。

推論結果を描画します。

from PIL import Image
import numpy as np
import io

num_classes = 21
mask = np.array(Image.open(io.BytesIO(return_img)))
plt.imshow(mask, vmin=0, vmax=num_classes-1, cmap='jet')
plt.show()

次にacceptタイプがプロトコルバッファ形式として推論を試します。

%%time
# 推論用に画像サイズを変更します
im = PIL.Image.open(filename)
im.thumbnail([800,600],PIL.Image.ANTIALIAS)
im.save(filename, "JPEG")
with open(filename, 'rb') as image:
    img = image.read()
    img = bytearray(img)
    
ss_predictor.content_type = 'image/jpeg'
ss_predictor.accept = 'application/x-protobuf'
results = ss_predictor.predict(img)

推論結果はプロトコルバッファ形式のRecordIOなのでMXNetの処理を使って扱いやすい形に変換します。

from sagemaker.amazon.record_pb2 import Record
import mxnet as mx

results_file = 'results.rec'
with open(results_file, 'wb') as f:
    f.write(results)

rec = Record()
recordio = mx.recordio.MXRecordIO(results_file, 'r')
protobuf = rec.ParseFromString(recordio.read())

中には2種類のデータが入っています。出力結果の形(shape)とピクセル毎の各クラスの確率データ(target)です。targetはそのままだと１次元の数値データなので、shapeに従って適切な配列形式に変換します。

values = list(rec.features["target"].float32_tensor.values)
shape = list(rec.features["shape"].int32_tensor.values)
shape = np.squeeze(shape)
mask = np.reshape(np.array(values), shape)
mask = np.squeeze(mask, axis=0)

targetデータを見てみます。データ量が多いので先頭五個だけ表示します。

values[:5]

shapeデータも見てみます。

shape

ピクセル毎に確率の最も高いクラスを選び、描画してみます。

pred_map = np.argmax(mask, axis=0)
num_classes = 21
plt.imshow(pred_map, vmin=0, vmax=num_classes-1, cmap='jet')
plt.show()

acceptタイプがpng画像と同じような感じで描画されました! 単にラベル分けされた結果が欲しい場合はacceptタイプとしてpngを使い、より詳細なデータが欲しい場合にはacceptタイプとしてprotobufを使うのが良さそうです。

エンドポイントの削除

最後に無駄な課金が発生しないようにエンドポイントを削除します。

sagemaker.Session().delete_endpoint(ss_predictor.endpoint)

さいごに

今回は 「クラスメソッド Amazon SageMaker Advent Calendar」の1日目として、 「Amazon SageMakerの組み込みアルゴリズムのセマンティックセグメンテーション」についてお送りしました。セマンティックセグメンテーションは物体検出や画像分類と同様、機械学習とマッチした問題で応用範囲も広いかと思いますし、それらに比べてさらに細かい位置の検出、ピクセル単位での検出が可能となります。教師データの作成は非常に大変ですが、用途によってはとても有用なアルゴリズムとなるのではないでしょうか。

最後までお読みくださり、ありがとうございました〜！
明日は、yoshimより「TensorFlow hub」に公開されているモデルをデプロイしてみる」についてお送りします。お楽しみに〜!