SageMaker Training Job でカスタムコンテナを用いてモデルを学習する

Training Job 概要

SageMaker Training Job では、Training Job というものを発行することで機械学習モデルを学習できます。^[1]
発行時にいろいろな設定を行うのですが、基本的には以下の 3 つの設定をするだけです。

• 学習で使用するコンテナイメージ (ECR)
• 学習データセットやモデル、コンフィグファイルが置いてある S3 パス
• 学習後のモデル出力用の S3 パス

これらの項目を設定してジョブを発行すると、

• インスタンスの立ち上げ
• S3 から入力用データのコピー
• コンテナの起動 (コンテナ内の /opt/ml/ 配下に入出力用ディレクトリがマウントされる)
• 学習コード実行
• S3 への出力データのアップロード

を SageMaker が自動で行ってくれます。^[2]

ユーザーが行うこと

上記のような学習を行うにあたって以下の 4 つをユーザー側が基本的に行わなければなりません。

• S3 バケットの作成と、入力用データのアップロード
• Dockerfile の作成
• 学習用コードの作成
• Training Job の設定の入力と発行

以下では、Dockerfile と学習用コードの解説を行っていきますが、入出力データを置くディレクトリは SageMaker が指定する /opt/ml/ 配下であることに注意してください。

/opt/ml/ 配下のディレクトリ構造

/opt/ml/ は SageMaker が S3 とのデータのやり取りで使用するディレクトリです。
コンテナ起動時にマウントされるディレクトリのため /opt/ml/ ディレクトリをユーザーが作成する必要はありません。
ツリーは主にこのようになっています。^[3]

opt
└── ml
    ├── input
    │   ├── config
    |   |   ├── init-config.json
    |   |   ├── checkpointconfig.json
    │   │   ├── hyperparameters.json
    │   │   ├── inputdataconfig.json
    │   │   └── resourceconfig.json
    │   └── data
    │       └── <channel_name>
    │           └── <input data>
    ├── model
    └── output

ツリーの説明は以下になります。

パス	説明
/opt/ml/input/config/	ジョブ発行時に行った設定を記述したファイルが置かれます。
/opt/ml/data/	ここにはジョブ発行時に設定した S3 のパスからデータがコピーされます。 S3 に存在するパスを複数指定することができ、 channel_name は、それぞれのパスに対応し、任意の名前を付けることができます。 チャンネル名の例：学習用のデータ : train 、モデルのパラメーター : model 、モデルのコンフィグ : config
/opt/ml/model/	ここに学習後のモデルを置いておくと、自動で S3 に model.tar.gz という名前でアップロードされます。
/opt/ml/output/	ここには学習失敗時のログを出力します。

Dockerfileについて

ユーザーが作成したファイルなどは基本的に、/opt/ 配下にディレクトリを作成して配置します。
今回は、/opt/program/train/train.py に学習用コードを記述します。
コンテナのエントリポイントを以下のように設定すると、コンテナ起動時に学習が走ります。

ENTRYPOINT ["python", "/opt/program/train/train.py"]

デフォルトでは、docker run 時に train コマンドが渡される形で、スクリプト /opt/program/train が実行されるようです。

今回は以下のようなDockerfileからイメージを作成します。

Dockerfile

FROM ubuntu:20.04

ENV DEBIAN_FRONTEND=noninteractive

RUN apt-get update && \
    apt-get -y install --no-install-recommends \
    wget \
    git \
    build-essential \
    ca-certificates \
    curl \
    libssl-dev \
    libffi-dev \
    python3.9 \
    python3.9-dev \
    libgl1-mesa-glx \
    libglib2.0-0 \
    python3-distutils \
    && rm -rf /var/lib/apt/lists/* \
    && curl -O https://bootstrap.pypa.io/get-pip.py \
    && python3.9 get-pip.py

RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.9 1
RUN update-alternatives --install /usr/local/bin/pip pip /usr/local/bin/pip3 1

ENV PROGRAM_DIR=/opt/program
WORKDIR $PROGRAM_DIR

# clone and install model dependencies
RUN git clone https://github.com/backprop64/DAMM && \
    pip --no-cache-dir install -r DAMM/requirements-gpu.txt
RUN cp -r DAMM/DAMM /usr/local/lib/python3.9/dist-packages/

ENV PYTHONUNBUFFERED=TRUE
ENV PYTHONDONTWRITEBYTECODE=TRUE

ENV PATH="/opt/program:${PATH}"

COPY train.py /opt/program/train/train.py
RUN chmod +x /opt/program/train/
RUN chmod +x /opt/program/train/train.py

ENTRYPOINT ["python", "/opt/program/train/train.py"]

CMD []

学習するモデルについて

今回は、DAMM というネズミ検知モデルをファインチューニングします。
pytorch を使用するため、上記の Dockerfile では、ネズミ検知モデルをクローンしてきてそれに従う依存関係をインストールしています。

学習用コード

/opt/ml/input/ からデータをとってきて、
/out/ml/output/model/ に学習データを出力する設定をしています。

train.py

#!/usr/bin/env python
from __future__ import print_function

import os
import sys
import json
import traceback
from pathlib import Path
from collections import namedtuple

import cv2

from DAMM.detection import Detector

def get_paths():
    ml_prefix = Path('/opt/ml/')
    param_path = ml_prefix / 'input' / 'config' / 'hyperparameters.json'
    input_train_dir = ml_prefix / 'input' / 'data' / 'train' / 'my_detection_dataset'
    input_model_path = ml_prefix / 'input' / 'data' / 'model' / 'model' / 'model_final.pth'
    input_model_config_path = ml_prefix / 'input' / 'data' / 'model' / 'model' / 'config.yaml'
    model_output_dir = ml_prefix / 'model'
    log_output_dir = ml_prefix / 'output'

    Paths = namedtuple('Paths', ['param_path', 'input_train_dir', 'input_model_path', 'input_model_config_path', 'model_output_dir', 'log_output_dir'])
    paths = Paths(param_path, input_train_dir, input_model_path, input_model_config_path, model_output_dir, log_output_dir)

    return paths


def run_inference():
    print_tree('/opt/ml/')
    paths = get_paths()

    print('Starting the inference.')
    try:
        # load hyperparameters but not used
        with paths.param_path.open() as tc:
            training_params = json.load(tc)

        # load DAMM detector
        damm_detector = Detector(
            cfg_path= str(paths.input_model_config_path),
            model_path= str(paths.input_model_path),
            output_dir=str(paths.model_output_dir),
        )

        # fine tune detector
        dataset_metadata_path = str(paths.input_train_dir / 'metadata.json')
        damm_detector.train_detector(dataset_metadata_path)

    except Exception as e:
        trc = traceback.format_exc()
        with open(str(paths.log_output_dir / 'failure'), 'w') as s:
            s.write('Exception during training: ' + str(e) + '\n' + trc)

        print('Exception during training: ' + str(e) + '\n' + trc,
              file=sys.stderr)
        sys.exit(255)


if __name__ == '__main__':
    run_inference()

Training Job の発行

ブラウザ上でポチポチするときの手順
ジョブ設定画面にて

コンテナ設定

インスタンス設定

今回は GPU を用いた学習を行うため、GPU の乗ったインスタンスタイプを選択します。

入力パス設定

ファインチューニング用の学習データを /opt/ml/input/data/train/ に、
ファインチューニング用のモデルを /opt/ml/input/data/model/ に配置しています。
S3 の mouse-training-job/input/data/train/ 直下のファイル、ディレクトリは直接 /opt/ml/input/data/train/ 下に配置されるので注意

出力パス設定

出力

ジョブ発行後、無事にモデルが出力されました！

まとめ

入出力にだけ気を付ければ、後はかなり自由に学習が行えるみたいでとても便利な印象。
今後は AWS Lambda からのジョブ発行や、学習後のモデルを S3 イベントで検知して推論などにつなげていきたいですね！

脚注

1. https://docs.aws.amazon.com/sagemaker/latest/dg/how-it-works-training.html

2. https://docs.aws.amazon.com/sagemaker/latest/dg/docker-containers.html

3. https://zenn.dev/kazuhito/books/4bfc0e4e39752c/viewer/56709b