MXNetのスクリプトでハイパーパラメータチューニングをする:Amazon SageMaker Advent Calendar 2018

それぞれ、順に見ていきましょう。

3-1.(if name=='main':)配下に色々と処理を記述する。

「1.3.0」以降では、「if name=='main':」配下に下記の3つの処理を実装する必要があります。

First, add a main guard (if name == 'main':). The code executed from your main guard needs to:
1.Set hyperparameters and directory locations
2.Initiate training
3.Save the model

引用:updating-your-mxnet-training-script

つまり、各種パラメータの引き渡し、学習、パラメータの保存、を記述する必要があるということですね。
学習処理については既存と変わりませんが、注意するのは「各種パラメータの引き渡し」と「パラメータの保存」の2点です。各パラメータの引き渡しについては後ほど説明しますが、argument parserを使うことが推奨されています。また、「パラメータの保存」については、下記のような処理を追加するだけでOKなパターンもありますが、場合によってはカスタムで処理を記述する必要があります。
パラメータを保存する関数を作成する際は「model」、「model_dir」の2つの引数を渡す必要がある点にご注意ください。

from sagemaker_mxnet_container.training_utils import save

if __name__ == '__main__':
    # arg parsing (shown above) goes here

    model = train(args.batch_size, args.epochs, args.learning_rate, args.train, args.test)
    save(args.model_dir, model)

引用:updating-your-mxnet-training-script

「model」、「model_dir」の2つの引数を渡す必要がある。

To provide your own save function, define a save function in your training script. The function should take two arguments:
model: This is the object that was returned from your train function. If your train function does not return an object, it will be None. You are free to return an object of any type from train, you do not have to return Module or Gluon API specific objects.
model_dir: This is the string path on the SageMaker training host where you save your model. Files created in this directory will be accessible in S3 after your SageMaker Training Job completes.

引用:writing-a-custom-save-function

3-2.各パラメータの引き渡し

今回参考にしたチュートリアルでは、MXNetのフレームワークバージョンが「1.2.0」だったために、エントリーポイントファイルに、「hyperparameters」引数に関する記述をしなくてもハイパーパラメータを受け取ることができていました。
しかし、「1.3.0」以降では、引数を受け取る処理をエントリーポイントファイル内に明示的に記述する必要があります。

Hyperparameters are passed to your script as arguments and can be retrieved with an argparse.ArgumentParser instance. For example, a training script might start with the following:

引用:For versions 1.3 and higher

この記述方法については、argument parserを使うことが推奨されており、またupdating-your-mxnet-training-scriptに具体的な記述方法が紹介されています。

import argparse
import os

if __name__ == '__main__':

    parser = argparse.ArgumentParser()

    # hyperparameters sent by the client are passed as command-line arguments to the script.
    parser.add_argument('--epochs', type=int, default=10)
    parser.add_argument('--batch-size', type=int, default=100)
    parser.add_argument('--learning-rate', type=float, default=0.1)

    # input data and model directories
    parser.add_argument('--model-dir', type=str, default=os.environ['SM_MODEL_DIR'])
    parser.add_argument('--train', type=str, default=os.environ['SM_CHANNEL_TRAIN'])
    parser.add_argument('--test', type=str, default=os.environ['SM_CHANNEL_TEST'])

    args, _ = parser.parse_known_args()

このようにしてハイパーパラメータ等の変数を渡すことができます。
ここまでを「ハイパーパラメータをどのようにエントリーポイントファイルに渡しているか」といった観点から整理すると、エントリーポイントファイル実行時に「argparse」で明示的にハイパーパラメータを渡しているといったことがわかるかと思います。

最後に細かいのですが、エントリーポイントファイル実行時に参照できる「環境変数」についてはlist-of-provided-environment-variables-by-sagemaker-containersをご参照ください。
自分でMXNetでエントリーポイントファイルを記述する方には必須の知識となります。

4.「1.3.0」以降でのスクリプト例

ここまでの内容を反映して、チュートリアルと同じような処理を「MXNet　1.3.0」以降のバージョン用に書き直したスクリプトを紹介します。
スクリプト自体はエントリーポイントファイル,呼び出しファイルの通りですが、修正前との差分についてざっくり説明しようと思います。

4-1.エントリーポイントファイル

変数の受け取りや、パラメータの保存用にモジュールを追加しています。

import argparse
from sagemaker_mxnet_container.training_utils import save
import json

受け取ったハイパーパラメータも、以前は「hyperparameters」という辞書型の変数に受け取っていましたが、それぞれの引数として受け取っているので「train」関数に修正が入っています。

def train(current_host, train_dir, test_dir, hosts, num_cpus, num_gpus, optimizer, num_epoch, learning_rate, batch_size):

そして、最後に「if name =='main':」配下に「変数の受け取り」、「学習の実行」、「パラメータファイルの出力」を実行しています。
受け取るハイパーパラメータもそれぞれ別々に記述して受け取っています。

if __name__ =='__main__':
    parser = argparse.ArgumentParser()

    # hyperparameters sent by the client are passed as command-line arguments to the script.
    parser.add_argument('--optimizer', type=str)
    parser.add_argument('--num_epoch', type=int, default=100)
    parser.add_argument('--learning_rate', type=float, default=0.1)

    # input data and model directories
    parser.add_argument('--train', type=str, default=os.environ['SM_CHANNEL_TRAIN'])
    parser.add_argument('--test', type=str, default=os.environ['SM_CHANNEL_TEST']) 
    parser.add_argument('--model-dir', type=str, default=os.environ['SM_MODEL_DIR'])
    parser.add_argument('--batch_size', type=int)    

    # system_param
    parser.add_argument('--num_gpus', type=int, default=os.environ['SM_NUM_GPUS'])
    parser.add_argument('--num_cpus', type=int, default=os.environ['SM_NUM_CPUS'])
    parser.add_argument('--current_host', type=str, default=os.environ['SM_CURRENT_HOST'])
    parser.add_argument('--hosts', type=str, default=json.loads(os.environ['SM_HOSTS']))

    args, _ = parser.parse_known_args()

    model = train(args.current_host, args.train, args.test, args.hosts, args.num_cpus, args.num_gpus, args.optimizer, args.num_epoch,args.learning_rate,args.batch_size)
    save(args.model_dir, model)

4-2.呼び出しファイル

本筋とはずれてしまいますが、学習結果やソースコードを保存するS3パスを指定しています。 それ以外は特に変更ありません。

bucket = 's3://my-bucket'
prefix = 'mxnet_hpo'
model_artifacts_location = '{0}/{1}/model'.format(bucket,prefix)
custom_code_upload_location = '{0}/{1}/code'.format(bucket,prefix)

estimator = MXNet(# entry_point='mnist.py',
    entry_point='mnist_130.py', # after 「MXNet 1.3.0」
    role=role,
    train_instance_count=1,
    train_instance_type='ml.m4.xlarge',
    sagemaker_session=sagemaker.Session(),
    output_path=model_artifacts_location, # add!!!
    code_location=custom_code_upload_location, # add!
    base_job_name='DEMO-hpo-mxnet',
    framework_version='1.3.0',
    hyperparameters={'batch_size': 100}
)

5.まとめ

「MXNet 1.3.0」以降のフレームワークバージョンでのハイパーパラメータチューニング処理を実行するための記述内容について整理できました。

学習結果を推論用エンドポイントとして利用するためには、他にも「モデルをロードする関数」や「推論処理の関数」等を記述する必要があるのですが、とりあえず今回はここまでにしておこうかと思います。

MXNetを利用しようとしている方の参考になれば幸いです。