SageMakerの「パイプモード」を使ってTensorFlowの処理を高速化する:Amazon SageMaker Advent Calendar 2018

概要

こんにちは、データインテグレーション部のyoshimです。
この記事は「クラスメソッド Amazon SageMaker Advent Calendar」の4日目の記事となります。

目次

• 1.やること
• 2.「パイプモード」について
• 3.TensorFlowの自作コードで学習をする方法
• 4.まとめ

1.やること

SageMakerの「パイプモード」という機能を調べたことがなかったので調べてみました。
また、チュートリアルが公開されていたので、このチュートリアルの内容をベースに紹介していきたいと思います。

2.「パイプモード」について

まず、「パイプモード」とはなんぞや、というところから調べました。
こちらを見てみると、下記のような機能のようです。

• データセットを学習コンテナにダウンロードするのではなく、トレーニングインスタンスにS3から直接ストリーミングする。
• その結果、学習の開始が早まり、必要なディスク容量も小さくて済む

この記述だけを見てもシンプルに良さそうな機能です。
また「通常のファイル読み込みよりもS3からのストリーミングの方が読み込みが早く」、「S３をソースにしているので実質無制限な量のデータを処理」できます。どうやら、大規模なデータを処理する時にはパイプモードを使った方が良さそうですね。

下記は、AWSがパイプモードのメリットを示すために、通常の「ファイルモード」とのパフォーマンスを比較したものです。
ニューヨーク市のタクシートリップレコードのデータセットを使ってPCAを学習させた際のパフォーマンスだそうです。
引用: Amazon SageMaker アルゴリズムのパイプ入力モードを使用する

JOB開始時間、IOのいずれもパフォーマンスが向上しており、結果として合計の学習時間が削減されています。

3.TensorFlowの自作コードで学習をする方法

続いて、「TensorFlow」の自作コードで学習をする方法について説明します。
今回のチュートリアルで用意したスクリプトは下記の通りです。
(「pipemode.py」)

import numpy as np
import os
import tensorflow as tf
from sagemaker_tensorflow import PipeModeDataset

from tensorflow.contrib.data import map_and_batch

PREFETCH_SIZE = 10
BATCH_SIZE = 64
NUM_PARALLEL_BATCHES = 2
DIMENSION = 1024
EPOCHS = 1

def estimator_fn(run_config, params):
    column = tf.feature_column.numeric_column('data', shape=(DIMENSION, ))
    return tf.estimator.LinearClassifier(feature_columns=[column], config=run_config)

def train_input_fn(training_dir, params):
    """Returns input function that would feed the model during training"""
    return _input_fn('train')

def eval_input_fn(training_dir, params):
    """Returns input function that would feed the model during evaluation"""
    return _input_fn('eval')

def _input_fn(channel):
    """Returns a Dataset for reading from a SageMaker PipeMode channel."""
    features = {
        'data': tf.FixedLenFeature([], tf.string),
        'labels': tf.FixedLenFeature([], tf.int64),
    }

    def parse(record):
        parsed = tf.parse_single_example(record, features)
        return ({
            'data': tf.decode_raw(parsed['data'], tf.float64)
        }, parsed['labels'])

    ds = PipeModeDataset(channel)
    if EPOCHS > 1:
        ds = ds.repeat(EPOCHS)
    ds = ds.prefetch(PREFETCH_SIZE)
    ds = ds.apply(map_and_batch(parse, batch_size=BATCH_SIZE,
                                num_parallel_batches=NUM_PARALLEL_BATCHES))

    return ds

「estimator_fn」がモデルの定義を、「train_input_fn」、「eval_input_fn」がトレーニング用/評価用データを読み込む処理を記述しています。
「train_input_fn」、「eval_input_fn」はそれぞれ必須ですが、「SageMaker TensorFlow docker image」を使う場合は「estimator_fn」ではなく、「model_fn」、「keras_model_fn」を使うこともできます。
引用:sagemaker-python-sdkのREADME

• Exactly one of the following:
  • model_fn: defines the model that will be trained.
  • keras_model_fn: defines the tf.keras model that will be trained.
  • estimator_fn: defines the tf.estimator.Estimator that will train the model.
• train_input_fn: preprocess and load training data.
• eval_input_fn: preprocess and load evaluation data.

ちなみに、「estimator_fn」の基本骨格としては、下記のような関数を記述する必要があります。
(引用元にあったサンプルです)

def estimator_fn(run_config, hyperparameters):
  # Logic to the following:
  # 1. Defines the features columns that will be the input of the estimator
  # 2. Returns an instance of a ``tensorflow.estimator`` passing in, the input run_config in the
  #    constructor.

  # Defines the features columns that will be the input of the estimator
  feature_columns = [tf.feature_column.numeric_column(INPUT_TENSOR_NAME, shape=[4])]
  # Returns the instance of estimator.
  return tf.estimator.DNNClassifier(feature_columns=feature_columns,
                                      hidden_units=[10, 20, 10],
                                      n_classes=3,
                                      config=run_config)

引用:sagemaker-python-sdkのREADME

「train_input_fn」、「eval_input_fn」については、下記のような処理を記述する必要があります。

def train_input_fn(training_dir, hyperparameters):
  # Logic to the following:
  # 1. Reads the **training** dataset files located in training_dir
  # 2. Preprocess the dataset
  # 3. Return 1)  a dict of feature names to Tensors with
  # the corresponding feature data, and 2) a Tensor containing labels
  return features, labels

引用:sagemaker-python-sdkのREADME

また、「_input_fn」では「パイプモードでデータを読み込むための特別な処理」をしています。
具体的には、「PipeModeDataset」というデータとして変換する処理が必要になります。

詳細はsagemaker-tensorflow-extensionをご参照いただきたいのですが、ここでは「recordio」形式のファイルを「PipeModeDataset」形式に変換しています。
(2018年11月末時点では、こちらのドキュメントを見ると「TFRecord」、「RecordIO」、「txt」の3種類のファイルしか対応できていないようです)

def _input_fn(channel):
    """Returns a Dataset for reading from a SageMaker PipeMode channel."""
    features = {
        'data': tf.FixedLenFeature([], tf.string),
        'labels': tf.FixedLenFeature([], tf.int64),
    }

    def parse(record):
        parsed = tf.parse_single_example(record, features)
        return ({
            'data': tf.decode_raw(parsed['data'], tf.float64)
        }, parsed['labels'])

    ds = PipeModeDataset(channel) # 「PipeModeDataset」に変換
    if EPOCHS > 1:
        ds = ds.repeat(EPOCHS)
    ds = ds.prefetch(PREFETCH_SIZE)
    ds = ds.apply(map_and_batch(parse, batch_size=BATCH_SIZE,
                                num_parallel_batches=NUM_PARALLEL_BATCHES))

    return ds

SageMakerでTensorFlowの自作コードを利用する方法の詳細については、下記もご参照ください。
Amazon SageMakerでTensorFlowを使ってIris分類してみた
sagemaker-python-sdkのREADME

続いて、学習用のパラメータを設定します。
「input_mode」が「Pipe」になっているのがポイントです。

from sagemaker.tensorflow import TensorFlow

# https://docs.aws.amazon.com/ja_jp/sagemaker/latest/dg/tf-example1-train.html
tensorflow = TensorFlow(entry_point='pipemode.py',
                        role=role,
                        framework_version='1.11.0',
                        input_mode='Pipe', # 「Pipe」にすることでパイプモードに設定可能。「File」にするとファイルモード。
                        output_path=model_artifacts_location, # 学習した結果が格納されるS３パス
                        code_location=custom_code_upload_location, # 実行したスクリプトを保存するS3パス
                        train_instance_count=1,
                        training_steps=1000,
                        evaluation_steps=100,
                        train_instance_type='ml.c4.xlarge')

最後に利用するデータセットのS３パスを指定して学習を開始すればOKです!!!

%%time
import boto3

# use the region-specific sample data bucket
region = boto3.Session().region_name

train_data = 's3://sagemaker-sample-data-{}/tensorflow/pipe-mode/train'.format(region)
eval_data = 's3://sagemaker-sample-data-{}/tensorflow/pipe-mode/eval'.format(region)

tensorflow.fit({'train':train_data, 'eval':eval_data})

4.まとめ

パイプモードを使うことで、TensorFlowを使ったモデル開発を効率化することができそうですが、そのためには下記の対応をする必要があるようです。

• 1.データを「PipeModeDataset」形式に変換する（2018年11月末時点では、「TFRecord」、「RecordIO」、「txt」のいずれかからのみ可能）
• 2.エントリーポイントファイルの作成
• 3.学習パラメータ指定時に「input_mode」パラメータに「Pipe」を指定

学習にかかる時間、金額等のコスト削減を考えている方がいらっしゃったら是非「パイプモード」を試してみてください。