LoginSignup
1

More than 1 year has passed since last update.

posted at

updated at

Amazon SageMakerで深層学習画風変換モデルのハイパーパラメータチューニングに挑戦してみた

この記事は、NTTテクノクロス Advent Calendar 2020の13日目です。

はじめまして。NTTテクノクロスの堀江と申します。
今年の弊社アドベントカレンダー執筆陣の中では最年少(現在入社4年目)となります。どうぞよろしくお願いします!

業務では、パブリッククラウド上でのシステム設計や環境構築支援を担当しています。その関係もあって現在絶賛勉強中のAWSと、趣味で勉強している深層学習の2つを組み合わせた取り組みを本日は発表したいと思います。


目次

  • はじめに
    • 記事要約 - 調査の観点と成果
  • 用語説明
    • Neural Style Transferとは
    • Amazon SageMakerとは
  • 実施内容
    • 環境構築
    • 単発のトレーニングジョブを実行
    • ハイパーパラメータチューニングを実行
    • チューニング結果を比較
  • おわりに
  • Appendix

はじめに

TensorFlowで自作した画像生成深層学習モデル - Neural Style Transfer - のハイパーパラメータチューニングに、Amazon SageMakerを利用して挑戦してみました。
その結果得られた、Amazon SageMakerの使用方法に関するノウハウや、ハイパーパラメータチューニングの結果を本記事で展開していきます。


記事要約 - 調査の観点と成果

  • Amazon SageMakerのトレーニングジョブやハイパーパラメータチューニングジョブって、そもそもどういう仕組みで実行されるの?
    • 実際にトレーニングジョブとハイパーパラメータチューニングジョブを実行してみることで、データや処理の流れを理解できたので、自分なりに整理してみました。
  • MNIST分類モデルのようなシンプルな構造のモデルだけじゃなく、画像生成モデルのような複雑な構造のモデルもAmazon SageMakerでトレーニングできるの?
    • デフォルト設定のトレーニングジョブ実行時に遭遇したエラーと、その回避方法を記載しています。
    • 例えばtf.functionとkerasが混在し、tf.Modulでクラス化されたような複雑な構造のモデルでも、トレーニングを実行することが出来ました。
  • Neural Style Transferって、ハイパーパラメータチューニングの効果があるの?
    • ハイパーパラメータをチューニングすることで、生成される画像の質が向上することを確認できました。(少なくとも自分の主観では…)

用語説明

本題に入る前に、重要な用語2つについては簡単にですが説明させてもらいます。
なお、「深層学習とは」「AWSとは」といった基本の部分については本記事での解説を省略させてもらいますのでご容赦ください。


Neural Style Transferとは

Neural Style Transfer(以降、長いのでnstとも省略します)とは、深層学習を利用した画像生成技術の一つです。
下記の例のように、ある画像の画風(スタイル)で、任意の別の画像(コンテンツ)の画風を変換することができます。

  • 犬の写真をワシリー・カンディンスキーの画風(スタイル)で変換したデモgif動画1

nstのデモgif

  • nstのサンプル作品集2

nstで生成される画像の例

記事要約内では「自作の画像生成深層学習モデル」と誇張した表現を用いてしまいましたが…今回利用するnstモデルは、TensorFlowの公式サイトに載っているサンプルノートブックのコードを独自に(趣味で)リファクタリングしたものになります。数学的なアルゴリズムの詳細についても公式サイトに載っているので、興味のある方はそちらをご覧になってみてください。

今回使用するモデルのコード(一部を抜粋)
class NstEngine(tf.Module):
    def __init__(self, content_shape, args, name=None):
        """
        content_shape : コンテンツ画像のshape e.g. (1, 512, 512, 3)
        args : Namespaceオブジェクト。
        """
        super(NstEngine, self).__init__(name=name)

        # サンプノートブックの通り、VGG19を特徴量抽出器として利用する。
        self.content_layers = ['block5_conv2']
        self.style_layers = ['block1_conv1', 'block2_conv1', 'block3_conv1', 'block4_conv1', 'block5_conv1']
        vgg = tf.keras.applications.VGG19(include_top=False, weights='imagenet') 
        outputs = [
            vgg.get_layer(target_layer).output 
            for target_layer in self.content_layers + self.style_layers
        ]
        self.model = tf.keras.Model([vgg.input], outputs)
        self.model.trainable = False

        # tf.function内部では変数を宣言できないので、モデルの初期化時に
        # 変数も初期化しておく必要がある。
        self.content_image = tf.Variable(tf.zeros(content_shape), dtype=tf.float32)
        self.loss = tf.Variable(tf.zeros((1)), dtype=tf.float32)

        self.style_image = None
        self.content_image_org = None
        self.style_image_org = None
        self.content_target = None
        self.style_target = None


        self.epoch = int(args.EPOCH)
        # 以下、後々チューニング対象になるハイパーパラメータ達
        learning_rate = float(os.environ.get("SM_HP_LEARNING_RATE", args.LEARNING_RATE))
        self.optimizer = tf.keras.optimizers.Adam(
            learning_rate=learning_rate,
            beta_1=0.99,
            epsilon=0.1,
        )
        self.content_weights = float(os.environ.get("SM_HP_CONTENT_WEIGHTS", args.CONTENT_WEIGHTS))
        self.style_weights = float(os.environ.get("SM_HP_STYLE_WEIGHTS", args.STYLE_WEIGHTS))
        self.total_variation_weights = float(os.environ.get("SM_HP_TOTAL_VARIATION_WEIGHTS", args.TOTAL_VARIATION_WEIGHTS))


    @tf.function
    def fit(self, content, style, content_org):
        """
        外部から呼び出されるエントリーポイント
        args:
            - content : 更新対象のコンテンツ画像 shape : (1, height, width, 3)
            - style : スタイル画像 shape : (1, height, width, 3)
            - content_org : オリジナルのコンテンツ画像 shape : (1, height, width, 3)

        return : 
            - スタイル画像の画風で更新されたコンテンツ画像 shape : (1, height, width, 3)
            - loss値 shape : (1)
        """

        self.content_image_org = content_org
        self.style_image_org = style

        self.content_image.assign(content)
        self.style_image = style

        self.content_target = self.call(self.content_image_org)['content']
        self.style_target = self.call(self.style_image_org)['style']

        for e in tf.range(self.epoch):
            self.loss.assign([self.step()])

        return self.content_image, self.loss

"""
以降の関数はサンプルノートブックの処理を流用。
詳細については省略
"""
    @tf.function
    def step(self):
        with tf.GradientTape() as tape:
            outputs = self.call(self.content_image)
            loss = self._calc_style_content_loss(outputs)
            loss += self.total_variation_weights*self._total_variation_loss()

        grad = tape.gradient(loss, self.content_image)
        self.optimizer.apply_gradients([(grad, self.content_image)])
        self.content_image.assign(self._clip_0_1())

        return loss

    def _calc_style_content_loss(self, outputs):
        style_outputs = outputs['style']
        content_outputs = outputs['content']

        style_loss = tf.add_n([
            tf.reduce_mean((style_outputs[name] - self.style_target[name])**2)
            for name in style_outputs.keys()
        ])
        style_loss *= self.style_weights / len(self.style_layers)

        content_loss = tf.add_n([
            tf.reduce_mean((content_outputs[name] - self.content_target[name])**2)
            for name in content_outputs.keys()
        ])
        content_loss *= self.content_weights / len(self.content_layers)

        loss = style_loss + content_loss
        return loss

    def _total_variation_loss(self):
        x_deltas, y_deltas = self._high_pass_x_y()
        return tf.reduce_sum(tf.abs(x_deltas)) + tf.reduce_sum(tf.abs(y_deltas))

    def _clip_0_1(self):
        clipped = tf.clip_by_value(
            self.content_image,clip_value_min=0.0, clip_value_max=1.0
        )
        return clipped

    def _high_pass_x_y(self):
        x_var = self.content_image[:, :, 1:, :] - self.content_image[:, :, :-1, :]
        y_var = self.content_image[:, 1:, :, :] - self.content_image[:, :-1, :, :]
        return x_var, y_var


    def call(self, input_image):
        input_image = input_image * 255.
        image = tf.keras.applications.vgg19.preprocess_input(input_image)
        outputs = self.model(image)

        content_outputs = outputs[:len(self.content_layers)]
        style_outputs = outputs[len(self.content_layers):]

        style_matrix = self._calc_gram_matrix(style_outputs)

        style_dict = {
            name: output 
            for name, output in zip(self.style_layers, style_matrix)
        }
        content_dict = {
            name: output 
            for name, output in zip(self.content_layers, content_outputs)
        }

        return {'style' : style_dict, 'content' : content_dict}

    def _calc_gram_matrix(self, input_tensors):
        results = []
        for input_tensor in input_tensors:
            result = tf.linalg.einsum('bijc,bijd->bcd', input_tensor, input_tensor)
            input_shape = tf.shape(input_tensor)
            num_locations = tf.cast(input_shape[1] * input_shape[2], tf.float32)
            results.append(result / num_locations)
        return results

Amazon SageMakerとは

Amazon SageMaker は、すべての開発者やデータサイエンティストが機械学習 (ML) モデルを迅速に構築、トレーニング、デプロイできるようにする完全マネージド型サービスです。SageMaker は高品質モデルの開発を容易にするため、機械学習の各プロセスから負荷の大きな部分を取り除きます。

Amazon SageMakerの公式サイトより

Amazon SageMakerとは、その名の通りAWSが提供しているサービスの1つです。機械学習(特に深層学習)が必要とする高速かつ大規模な計算リソースを個人でも低コスト且つ気軽に利用することができます。
Amazon SageMakerのトレーニングジョブ環境はAWS側で管理してくれるため、EC2やECSを使用するよりも手軽にトレーニング環境を用意することがきで、かつ、コスト効率的に利用することが出来ます。(EC2やECSは仮想マシンが起動している時間で課金されるのに対し、SageMakerのトレーニング環境はトレーニングジョブが実行されている間のみ課金されます。)
Amazon SageMakerは機械学習モデルのトレーニング、チューニングに限らず、プロダクト開発者・研究者向けに様々な便利サービスを提供しています。しかし本記事ではあくまでもモデルのトレーニングとチューニング機能にのみフォーカスさせてもらいます。


実施内容

前置きが少々長くなりましたが…ここからが本題となります。nstモデルを実際にAmazon SageMakerでトレーニング、ハイパーパラメータチューニングしていきます。


環境構築

まずはSageMakerサービスの利用(クライアント)環境を、ローカルのMacBook(無印12インチ)上にDockerイメージとして構築していきます。
Amazon SageMakerでは、機械学習モデル開発に適した環境をAWS上に簡単に構築してくれるサービスも提供しています。必要なライブラリやJupyter環境がプレインストールされたノートブックインスタンスや、それに追加で更に統合開発環境的なユーティリティを備えたSageMaker Studio等です。
これらを使用する手もあるのですが、インスタンスに極力お金を掛けたくなかったのでオンプレ環境からもAmazon SageMakerのサービスを利用できるのか、今回ついでに確認したかったので、今回は使用しません。


Dockerイメージ

Dockerfile.dev
# tensorflow 2.3.0をベースイメージとして使用。
# デフォルトでtensorboardも利用可能。
# 面倒だったので、GPU向けのイメージのみビルドします。
# GPU利用不可能な環境でも、特に不具合なく起動してくれるので。
FROM tensorflow/tensorflow:2.3.1-gpu

USER root
ENV PYTHONPATH=/app/notebook

COPY ./src /app/notebook
COPY ./keras /root/.keras
COPY ./config /config

# pipでsagemakerのSDKをインストールすることで、SageMakerサービスが利用可能になる。
RUN pip install --no-cache-dir \
        matplotlib \
        Pillow==7.1.1 \
        boto3==1.14.44 \
        sagemaker==2.16.1 \
        jupyterlab

# 念のためAWS CLIもインストール
RUN curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip" \
    && unzip awscliv2.zip \
    && ./aws/install \
    && rm -fr awscliv2*

RUN chmod +x /config/entrypoint.sh
WORKDIR /app/notebook
# コンテナ起動と同時にjupyterlabも立ち上がるようにエントリーポイントを作成
ENTRYPOINT ["/config/entrypoint.sh"]

docker-composeスタック

docker-compose-cpu.yml
version: '2.4'

services:
  notebook:
    image: notebook:dev
    # runtime: nvidia
    container_name: notebook
    hostname: notebook
    build:
      context: ./notebook
      dockerfile: ./dockerfile/Dockerfile.dev

    environment:
      - PYTHONPATH=/app/notebook
      - AWS_REGION=ap-northeast-1
    volumes:
      - ./notebook/src:/app/notebook
      - ./notebook/config:/config
      - ./notebook/keras:/root/.keras
      - ~/.aws:/root/.aws # ホストマシン上のクレデンシャルを利用する
    ports:
      - "6006:6006" # tensorboard
      - "8888:8888" # jupyterlab

あとはイメージをビルドしてコンテナを起動し、jupyterlab(http://localhost:8888/)にアクセスすれば完了です。

$ docker-compose -f docker-compose-cpu.yml build
$ docker-compose -f docker-compose-cpu.yml up

CloudFormationスタック

AWS上に最低限用意する必要があるリソースは、モデルのトレーニング結果や訓練データを格納するためのS3バケットと、そのS3バケットにSageMakerがアクセスするためのIAMサービスロールのみです。

cloudformation.yml
AWSTemplateFormatVersion: 2010-09-09
Description: ---

Parameters: 
  BucketName:
    Description: a name for the bucket used by sagemaker
    Type: String
    Default: sagemaker-nst

Resources: 
  S3Bucket:
    Type: AWS::S3::Bucket
    Properties: 
      BucketName:
        Ref: BucketName


  SageMakerExecRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
        - Effect: Allow
          Principal:
            Service:
            - sagemaker.amazonaws.com
          Action:
          - sts:AssumeRole
      Path: /
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/AmazonSageMakerFullAccess
      Policies:
        - PolicyName: S3BucketAccessPolicy
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - "s3:GetObject"
                  - "s3:PutObject"
                  - "s3:DeleteObject"
                  - "s3:ListBucket"
                Resource:
                  Fn::Sub: 'arn:aws:s3:::${BucketName}'

Outputs:  
  S3Bucket4Train:
    Description: the s3 bucket used by sagemaker during training
    Value:
      Ref: BucketName
  SageMakerRole:
    Description: service role for sagemaker
    Value:
      Ref: SageMakerExecRole

単発のトレーニングジョブを実行

では、環境が構築できたので試しに1回、トレーニングジョブを実行してみます。
Amazon SageMakerの基本的なトレーニングジョブの流れは以下の通りです。

  1. ローカルのノートブック上で、モデルが定義されたトレーニングスクリプトを用意する
  2. ローカルのノートブック上でトレーニングの設定を定義する
  3. ローカルのノートブック上からトレーニングの実行をAmazon SageMakerにリクエストする
  4. ローカルのノートブック上から、トレーニングスクリプトをはじめとする、トレーニングに必要な資材がS3バケットにアップロードされる。
  5. AWS上でトレーニング環境(実態はコンテナ?)が起動し、必要なセットアップが為される。
  6. S3上のトレーニングスクリプトや訓練データがトレーニング環境にダウンロードされて、トレーニングスクリプトが実行される。この間トレーニングの進捗状況やスクリプト内でprintしたログはローカルのノートブック上やCloudWatchLogsに出力され続ける。
  7. 訓練中に生成されるデータ(TensorBoard用のログ等)や訓練が完了したモデルがトレーニング環境のファイルシステム上に出力される。
  8. トレーニングが完了したら、トレーニング環境のファイルシステム上に出力されたデータが、指定したS3バケットにアップロードされる。
  9. トレーニング環境が消滅する。

Amazon SageMakerの基本的なトレーニングジョブのフロー図


トレーニングジョブの具体例

  • TrainingJob.ipynb
import json
import boto3
import sagemaker
from sagemaker.tensorflow import TensorFlow
print("boto3 : ", boto3.__version__)
print("sagemaker : ", sagemaker.__version__)
# boto3 :  1.14.44
# sagemaker :  2.16.1
# utilities
get_s3_path = lambda *args: "s3://" + "/".join(args)
# シークレット情報(アカウントID等)をノートブック上に貼りたくないので、
# ファイルに予め記載しておき、それを読み込む。
with open("./secrets.json", "r", encoding="utf-8") as fp:
    secrets = json.load(fp)
role=secrets["RoleArn"]
s3_bucket=secrets["S3Bucket"]
print("S3 bucket for training : ", s3_bucket)
# S3 bucket for training :  sagemaker-nst
# AWSとのセッションを確立する
boto_session = boto3.Session(region_name="ap-northeast-1")
sess = sagemaker.Session(
    # リージョンを東京リージョンに指定
    boto_session=boto_session,
    # SageMkaerが使用するS3バケットを指定
    default_bucket=s3_bucket,
)
# トレーニング設定の定義に必要なパラメータを用意

# トレーニングで使用されるデータのダウンロード元、
# トレーニングで生成されるデータのアップロード先となるS3のパスを指定
s3_train_path = get_s3_path(sess.default_bucket(), "train")

# インスタンスタイプに"local"を指定すると、トレーニングジョブの実行環境のDockerイメージがpullされ、
# ローカル上でコンテナが起動しトレーニングが実行される
# CPU、GPUともに安い価格帯のインスタンスを今回は使用する。
instance_types = {"CPU" : "ml.m5.large", "GPU" : "ml.g4dn.xlarge", "LOCAL" : "local"}

# トレーニングスクリプトに渡したいパラメータ
hyperparameters = {
    "EPOCH" : 50,
    "STEP" : 10,
    "MAX_IMAGE_SIZE" : 1024,
    "TB_BUCKET" : s3_train_path,
    "MAX_TRIAL" : 1,
}
# TensorFlowで書かれたモデルのトレーニングジョブ設定を定義
estimator = TensorFlow(
    # トレーニングスクリプト
    entry_point='train.py',
    # 指定したS3にSageMakerがアクセスするためのサービスロール
    role=role,
    # 起動されるインスタンス数
    instance_count=1,
    # 起動されるインスタンスのタイプ
    instance_type=instance_types["GPU"],
    # トレーニングジョブ実行環境内で使用されるtensorflowのバージョン
    framework_version='2.3.0',
    # トレーニングジョブ実行環境内で使用されるPythonのバージョン
    py_version='py37',
    # トレーニング中のデバッグライブラリの使用を無効化。理由は後述
    debugger_hook_config=False,
    # JSON形式で指定しtパラメータは、トレーニングジョブ実行時にコマンドライン引数としてトレーニングスクリプトに渡される。
    hyperparameters=hyperparameters,
    # ユーザが指定したS3バケットやリージョンを使用したいので、カスタマイズしたセッション情報も渡す
    sagemaker_session=sess,
    # スポットインスタンスの使用を有効にする。nstモデルの場合、料金を約70%節約できる。
    use_spot_instances=True,
    max_run=3600,
    max_wait=3600,
)
# トレーニングジョブを実行
# 指定したS3バケットのパスから訓練用データがトレーニングジョブ実行環境上に自動ダウンロードされる。
# デフォルト(logs="All")だとログが大量に出すぎるのでここでは抑制する。
estimator.fit(
    s3_train_path,
    job_name="TestSingleTrainingJob",
    logs="None",
    wait=False,
)
  • トレーニングスクリプト - train.py(一部を抜粋)
train.py

def _parse_args():
    """
    SageMakerは諸々のパラメータをコマンドライン引数または環境変数としてトレーニングスクリプトに渡してくるので、parseargでパラメータを取得可能。

    return : Tuple(Namespace, List[str])
    """
    parser = argparse.ArgumentParser()

    # sagemakerが引数として渡してくるパラメータ
    parser.add_argument('--model_dir', type=str)
    parser.add_argument('--sm-model-dir', type=str, default=os.environ.get('SM_MODEL_DIR', "models"))
    parser.add_argument('--train', type=str, default=os.environ.get('SM_CHANNEL_TRAINING', "train"))
    parser.add_argument('--sm-output-dir', type=str, default=os.environ.get('SM_OUTPUT_DATA_DIR', "outputs"))
    parser.add_argument('--hosts', type=list, default=json.loads(os.environ.get('SM_HOSTS', "{}")))
    parser.add_argument('--current-host', type=str, default=os.environ.get('SM_CURRENT_HOST', socket.gethostname()))


    # ユーザが設定したパラメータ。後々チューニングする予定
    parser.add_argument('--CONTENT_WEIGHTS', type=float, default=10000)
    parser.add_argument('--STYLE_WEIGHTS', type=float, default=0.01)
    parser.add_argument('--TOTAL_VARIATION_WEIGHTS', type=float, default=30)
    parser.add_argument("--LEARNING_RATE", type=float, default=0.02)
    parser.add_argument("--STYLE_RESIZE_METHOD", type=str, default="original")

    # ユーザが設定したパラメータ
    parser.add_argument('--EPOCH', type=int, default=20)
    parser.add_argument('--STEP', type=int, default=25)
    parser.add_argument("--MAX_IMAGE_SIZE", type=int, default=512)
    parser.add_argument("--TB_BUCKET", type=str, default="")
    parser.add_argument("--MAX_TRIAL", type=int, default=25)

    return parser.parse_known_args()

if __name__ =='__main__':

    # parse arguments
    args, _ = _parse_args()
    print(args)

    """
    以降、長いので省略
    - 訓練データの前処理
    - Tensorboardのロギングを設定
    - nstを実行。
    - モデルやnstされた画像の保存処理
    """

    return

トレーニングジョブを実行して5分ほど経過したら、AWS CLIでトレーニングジョブのステータスを確認してみます。

トレーニングジョブのステータスを確認
$ aws sagemaker describe-training-job --training-job-name TestSingleTrainingJob | jq .
{
  "TrainingJobName": "TestSingleTrainingJob",
  "TrainingJobArn": "arn:aws:sagemaker:ap-northeast-1:XXXXXXXXXXXX:training-job/testsingletrainingjob",
  "ModelArtifacts": {
    "S3ModelArtifacts": "s3://sagemaker-nst/TestSingleTrainingJob/output/model.tar.gz"
  },
  "TrainingJobStatus": "Completed",
  "SecondaryStatus": "Completed",
  "HyperParameters": {
    "EPOCH": "50",
    "MAX_IMAGE_SIZE": "1024",
    "MAX_TRIAL": "1",
    "STEP": "10",
    "TB_BUCKET": "\"s3://sagemaker-nst/train\"",
    "model_dir": "\"s3://sagemaker-nst/tensorflow-training-2020-11-14-03-52-18-413/model\"",
    "sagemaker_container_log_level": "20",
    "sagemaker_job_name": "\"TestSingleTrainingJob\"",
    "sagemaker_program": "\"train.py\"",
    "sagemaker_region": "\"ap-northeast-1\"",
    "sagemaker_submit_directory": "\"s3://sagemaker-nst/TestSingleTrainingJob/source/sourcedir.tar.gz\""
  },
  ...
}

"TrainingJobStatus"が"Completed"と表示されいていることから、設定通りにジョブが実行され無事成功したことが確認できます。
最後に、トレーニングスクリプト内でS3に出力されたログをTensorboardで確認してみます。

$ tensorboard  --host 0.0.0.0 --port 6006 --logdir s3://sagemaker-nst/train/tensorflow-training-2020-11-14-03-52-18-413
2020-11-14 04:30:44.993654: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcudart.so.10.1
TensorBoard 2.3.0 at http://0.0.0.0:6006/ (Press CTRL+C to quit)
  • TensorBoardでの、単発トレーニングジョブの確認結果3

トレーニングジョブ中に出力されたログの確認

無事、nstモデルのトレーニングジョッブ実行に成功しました。
余談ですが、TensorBoardのログディレクトリとしてS3バケット上のオブジェクトパスをダイレクトに指定できるということを、今回初めて知りました…。非常に便利ですね笑


トレーニングジョブ実行時に遭遇したエラーと、その回避方法

先述の具体例では、トレーニングジョブの定義時にdebugger_hook_config=False,という引数を指定しました(デフォルトではTrue)。この引数は、sagemaker-debuggerというライブラリを、Amazon SageMakerがトレーニングジョブ実行時に使用するかどうかを指定するためにあります。sagemaker-debugger自体の機能は自分も今回深掘りしていないのであまり把握していませんが、トレーニング中のモデルに発生する様々な問題(勾配消失等)を検出してくれるデバッグライブラリのようです。
私が今回使用しているnstモデルのコードでは、sagemaker-debuggerが有効になっているとトレーニング中にエラーを引き起こしてしまいます。

事象再現のために使用したスクリプト
import tensorflow as tf

class issueReproducer(tf.Module):

    def __init__(self, n_unit):
        """
        n_unit : int
        """
        self.variable = tf.Variable(tf.zeros((1, n_unit), dtype=tf.float32))
        self.l1 = tf.keras.layers.Dense(n_unit)
        self.optimizer = tf.optimizers.Adam()

    @tf.function
    def fit(self, tensor):
        """
        tensor : some tensor of shape : (1, n_unit)
        """
        with tf.GradientTape() as tape:
            output = self.l1(self.variable)
            loss = tf.reduce_sum(output - self.variable)
        grad = tape.gradient(loss, self.variable)
        self.optimizer.apply_gradients([(grad, self.variable)])

        return self.variable


if __name__ == "__main__":

    model = issueReproducer(5)
    tensor = tf.constant([[1,2,3,4,5]], dtype=tf.float32)

    variable = model.fit(tensor)

    print("Returned variable : {}".format(variable))

事象再現時のTraceback
Traceback (most recent call last):
  File "issue_reproducer.py", line 34, in <module>
    variable = model.fit(tensor)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 780, in __call__
    result = self._call(*args, **kwds)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 823, in _call
    self._initialize(args, kwds, add_initializers_to=initializers)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 697, in _initialize
    *args, **kwds))
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 2855, in _get_concrete_function_internal_garbage_collected
    graph_function, _, _ = self._maybe_define_function(args, kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 3213, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 3075, in _create_graph_function
    capture_by_value=self._capture_by_value),
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py", line 986, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 600, in wrapped_fn
    return weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 3735, in bound_method_wrapper
    return wrapped_fn(*args, **kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py", line 973, in wrapper
    raise e.ag_error_metadata.to_exception(e)
tensorflow.python.framework.errors_impl.OperatorNotAllowedInGraphError: in user code:

    issue_reproducer.py:23 fit  *
        grad = tape.gradient(loss, self.variable)
    /usr/local/lib/python3.7/site-packages/smdebug/tensorflow/keras.py:956 run  **
        (not grads or not vars)
    /usr/local/lib/python3.7/site-packages/tensorflow/python/framework/ops.py:877 __bool__
        self._disallow_bool_casting()
    /usr/local/lib/python3.7/site-packages/tensorflow/python/framework/ops.py:487 _disallow_bool_casting
        "using a `tf.Tensor` as a Python `bool`")
    /usr/local/lib/python3.7/site-packages/tensorflow/python/framework/ops.py:474 _disallow_when_autograph_enabled
        " indicate you are trying to use an unsupported feature.".format(task))

    OperatorNotAllowedInGraphError: using a `tf.Tensor` as a Python `bool` is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.

sagemaker-debuggerのドキュメントを確認すると、以下のような記載がありました。TensorFlowの一部の部品にはまだ対応しておらず、トレーニングスクリプトのコードを変更することなく利用できる範囲はまだ限られているようです。

* Debugger with zero script change is partially available for these TensorFlow versions. The inputs, outputs, gradients, and layers built-in collections are currently not available for these TensorFlow versions.

2020年11月27日時点のAmazon SageMaker DebuggerのDeveloper Guideより

Tracebackを確認してみると、確かに"not available"の例として挙げられているgrad = tape.gradient(loss, self.variable)の部分でエラーになっていますね。

本エラーの回避方法としてひとまず有効だったのが、前述したとおりトレーニングジョブの定義時にdebugger_hook_config=Falseを指定してsagemaker-debugger自体を無効化することでした。
スクリプトの書き方を工夫すればsagemaker-debuggerがONのままでも大丈夫なのか分かりませんが、いったん以降はOFFにしたままで進めていきます。


ハイパーパラメータチューニングを実行

単発のトレーニングジョブが無事(?)成功したので、次は本題となるハイパーパラメータチューニングを実践してみます。
前述したトレーニングジョブの手順に幾つかの設定を追加するだけで、ハイパーパラメータチューニングを実施できます。

  1. ローカルのノートブック上でトレーニングスクリプトを用意する。
  2. ローカルのノートブック上でトレーニングの設定を定義する。
  3. ローカルのノートブック上でハイパーパラメータチューニングの設定を定義する。
  4. ローカルのノートブック上から、Amazon SageMakerに対してハイパーパラメータチューニングの実行をリクエストする。
  5. 設定したハイパーパラメータの組み合わせごとに、指定した回数だけトレーニングジョブが実行される。
  6. 各トレーニングジョブからCloudWatchLogsに出力される目標メトリクスの値を集計・比較する。
  7. 実行されたトレーニングのうち、最善の性能を発揮したハイパーパラメータの値の組み合わせが確認できる。

Amazon SageMakerの基本的なハイパーパラメータチューニングのフロー図


ハイパーパラメータチューニングの具体例

  • HyperParameterTuning.ipynb
import boto3
import sagemaker
from sagemaker.tensorflow import TensorFlow
from sagemaker.tuner import IntegerParameter, CategoricalParameter, ContinuousParameter, HyperparameterTuner
import json
print("boto3 : ", boto3.__version__)
print("sagemaker : ", sagemaker.__version__)

# トレーニングジョブの設定定義
# この処理は単発のトレーニングジョブ実行時と同様なので省略
...

estimator = TensorFlow(
  ...
)
# チューニングするパラメータと、取り得るパラメータの範囲を定義
# 今回はloss値の計算に使用される3種類のパラメータと、
# Adamオプティマイザの学習率をチューニング対象にします。
hyperparameter_ranges = {
    "CONTENT_WEIGHTS" : ContinuousParameter(
        min_value=5000, max_value=15000,
    ),
    "STYLE_WEIGHTS" : ContinuousParameter(
        min_value=0.001, max_value=0.1,
    ),
    "TOTAL_VARIATION_WEIGHTS" : ContinuousParameter(
        min_value=10, max_value=50,
    ),
    "LEARNING_RATE" : ContinuousParameter(
        min_value=0.01, max_value=0.1,
    ),
}
# モデルの性能を比較するための目標メトリクスを定義
# loss値が最小になるハイパーパラメータの組み合わせを今回は探ってみる
objective_metric_name = 'loss'
objective_type = 'Minimize'
metric_definitions = [
    {
        'Name': 'loss',
        # トレーニングスクリプト内でprint()標準出力された
        # メッセージはCloudWatchLogsにログ出力される。
        # そのメッセージのうち、"FinalMeanLoss=([0-9\\.]+)"のパターンに合致する
        # 数値を目標メトリクスとして収集し、トレーニングジョブ間で比較する。
        'Regex': 'FinalMeanLoss=([0-9\\.]+)',
    }
]
# ハイパーパラメータチューニングジョブ設定を定義
tuner = HyperparameterTuner(
    # トレーニングジョブの定義
    estimator=estimator,
    # チューニング対象のハイパーパラメータ
    hyperparameter_ranges=hyperparameter_ranges,
    # 目標メトリクスの定義(今回はloss値を最小にする)
    objective_type=objective_type,
    objective_metric_name=objective_metric_name,
    metric_definitions=metric_definitions,
    # トータルで実行されるトレーニングジョブの上限数
    max_jobs=30,
    # 並列実行されるジョブの上限
    # ml.g4dn.xlargeインスタンスの同時起動可能上限数は2
    max_parallel_jobs=2,
)
# ハイパーパラメータチューニングの実行
tuner.fit(s3_train_loc, wait=False)
  • トレーニングスクリプト - train.py(一部を抜粋)
train.py

def _parse_args():

    parser = argparse.ArgumentParser()

    # sagemakerが引数として渡してくるパラメータ
    ...

    # 今回チューニングされるパラメータ
    # ハイパーパラメータチューニング時も、実際の値はコマンドライン引数として
    # トレーニングスクリプトに渡されるので、スクリプト内の処理を改修する必要無し。
    parser.add_argument('--CONTENT_WEIGHTS', type=float, default=10000)
    parser.add_argument('--STYLE_WEIGHTS', type=float, default=0.01)
    parser.add_argument('--TOTAL_VARIATION_WEIGHTS', type=float, default=30)
    parser.add_argument("--LEARNING_RATE", type=float, default=0.02)
    parser.add_argument("--STYLE_RESIZE_METHOD", type=str, default="original")

    ...

    return parser.parse_known_args()

if __name__ =='__main__':

    # parse arguments
    args, _ = _parse_args()
    print(args)

    ...

    print("FinalMeanLoss={}".format(loss_list.mean()))

チューニングジョブが完了したら、AWS CLIでチューニングの結果を確認してみます。

$ aws sagemaker describe-hyper-parameter-tuning-job --hyper-parameter-tuning-job-name tensorflow-training-201118-1130 | jq .
{
  ...
  "BestTrainingJob": {
    "TrainingJobName": "tensorflow-training-201118-1130-027-6b526460",
    "TrainingJobArn": "arn:aws:sagemaker:ap-northeast-1:XXXXXXXXXXXX:training-job/tensorflow-training-201118-1130-027-6b526460",
    "CreationTime": "2020-11-19T02:07:22+09:00",
    "TrainingStartTime": "2020-11-19T02:10:33+09:00",
    "TrainingEndTime": "2020-11-19T02:34:11+09:00",
    "TrainingJobStatus": "Completed",
    "TunedHyperParameters": {
      "CONTENT_WEIGHTS": "5277.5099011496795",
      "LEARNING_RATE": "0.03725186282831269",
      "STYLE_WEIGHTS": "0.0012957731765036497",
      "TOTAL_VARIATION_WEIGHTS": "10.0"
    },
    "FinalHyperParameterTuningJobObjectiveMetric": {
      "MetricName": "loss",
      "Value": 2388418.5
    },
    "ObjectiveStatus": "Succeeded"
  }
}

コード上からも直接、最善のハイパーパラメータの組み合わせを取得してみましょう。

HyperParameterTuning.ipynbの続き
# 最善の結果を出したモデルの、ハイパーパラメータの組み合わせを確認する。
from IPython.display import display
best_hyperparameters = tuner.best_estimator().hyperparameters()
display(best_hyperparameters)
# 2020-11-18 17:34:11 Starting - Preparing the instances for training
# 2020-11-18 17:34:11 Downloading - Downloading input data
# 2020-11-18 17:34:11 Training - Training image download completed. Training in progress.
# 2020-11-18 17:34:11 Uploading - Uploading generated training model
# 2020-11-18 17:34:11 Completed - Training job completed
# {'CONTENT_WEIGHTS': '5277.5099011496795',
#  'EPOCH': '50',
#  'LEARNING_RATE': '0.03725186282831269',
#  'MAX_IMAGE_SIZE': '1024',
#  'MAX_TRIAL': '9',
#  'STEP': '10',
#  'STYLE_WEIGHTS': '0.0012957731765036497',
#  'TB_BUCKET': '"s3://sagemaker-nst/train"',
#  'TOTAL_VARIATION_WEIGHTS': '10.0',
#  '_tuning_objective_metric': '"loss"',
#  'sagemaker_container_log_level': '20',
#  'sagemaker_estimator_class_name': '"TensorFlow"',
#  'sagemaker_estimator_module': '"sagemaker.tensorflow.estimator"',
#  'sagemaker_job_name': '"tensorflow-training-2020-11-18-11-30-35-845"',
#  'sagemaker_program': '"train.py"',
#  'sagemaker_region': '"ap-northeast-1"',
#  'sagemaker_submit_directory': '"s3://sagemaker-nst/tensorflow-training-2020-11-18-11-30-35-845/source/sourcedir.tar.gz"',
#  'model_dir': '"s3://sagemaker-nst/tensorflow-training-2020-11-18-11-26-25-821/model"'}

ハイパーパラメータチューニングジョブについても、nstモデルを対象に実施できることが確認出来ました。

因みに今回のハイパーパラメータチューニングジョブ内で、GPU利用可能なインスタンスタイプ(ml.g4dn.xlarge : 0.994USD/h)で合計30トレーニング(2並列)のトレーニングジョブが実行されたわけですが、その課金額はざっと以下の通りです。
※スポットインスタンスの使用を有効化しているため、課金額は(課金対象時間(実際にトレーニングが実行されている時間) * インスタンスの時間単価)で計算されます。今回使用しているnstモデルとデータ量であれば、60~70%程度節約されています。
※今回は2並列で実行したので、体感の経過時間は総実行時間 / 2です。

総実行時間 課金対象トレーニング時間 割引率 課金額
12時間8分 4時間16分 65% $4.24(約440円)

Amazon SageMakerを題材に勉強してみようと思い立った時は、文字通りの学習コストが最終的に幾ら位になるのかとビクビクしておりましたが、ス○バのカフェモカ一杯程度の金額に収まってくれて今はホッとしています笑。


チューニング結果を比較

ハイパーパラメータチューニンングが無事成功したので、それぞれのモデルのloss値や実際に生成された画像を比較してみたいと思います。
最善、最悪、デフォルトのハイパーパラメータの組み合わせと、その際のlossは以下の通りです。
なお、ここでの「デフォルト」とはTensorflowのサンプルコードに記載されていた設定値のことを指します。

PATTERN LOSS CONTENT_WEIGHT STYLE_WEIGHT TOTAL_VARIATION_WEIGHT LEARNING_RATE
最悪 20,893,880 5119.428 0.0752 12.171 0.078
デフォルト 10,061,276 10000 0.01 30 0.02
最善 2,388,418.5 5277.509 0.001 10 0.037

最善の結果はデフォルトのそれと比較しても4分の1以下、最悪のそれに対しては9分の1程もloss値が小さい結果となりました。

では、それぞれのトレーニングジョブで生成した9種類の画像をそれぞれ比較してみます。
※左から順に、「コンテンツ画像」、「スタイル画像」、「最悪のモデルの生成画像」、「デフォルトのモデルの生成画像」、「最善のモデルの生成画像」

最悪・デフォルト・最善のモデルの生成画像比較結果

最悪のモデルが生成した画像は論外として、デフォルトのモデルが生成した画像と比較しても、最善のモデルが生成した画像はスタイル画像の画風をしっかり反映しつつ、元の輪郭がよりハッキリと残っているように見えます。(少なくとも個人的には…)

最後に、モデルの汎化性能(めいたもの)も確認しておく意味で、新しい画像セットを対象に最良のモデルでnstを実施してみます。

  • 検証画像セットに対するnst実行結果4

検証画像セットに対する最良モデルのnst結果

良い感じにの画像が生成できている気がします(自己暗示)!
nstモデルに対するハイパーパラメータチューニングの効果があったと言えるのではないでしょうか!


おわりに

以上で、Amazon SageMakerを使用したnstモデルのハイパーパラメータチューニングの挑戦は終了となります。
見切り発車的に思い切って初立候補した今回のアドベントカレンダーでしたが、何とか形になってホッとしております…。
今まで使用経験が全く無かったAWSサービスの一つ(Amazon SageMaker)を集中的に学べたのは大変有意義でしたし、ハイパーパラメータチューニングの結果として、nstモデルが実際に生成する画像の質の改善を確認できたのは、正直自分でも驚きでした笑。

最後に。この記事で私が整理した図表、記載したコード、報告したエラーのいずれかが、少しでも誰かの役に立ってくれたり、刺激になってくれたりしたら大変嬉しいです。
みなさま。よいクリスマスを。

NTTテクノクロス Advent Calendar 2020、明日は@geek_duckさんです。お楽しみに!


Appendix


Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
What you can do with signing up
1