はじめに

こんにちは、研究開発部の石井です。

本エントリーでは、弊社の OCR エンジン「NineOCR」の開発フローに SageMaker Training を導入した話を紹介します。

NineOCR とは

NineOCR は Sansan が独自に開発した名刺特化の OCR エンジンです。

名刺をデータ化するフローの中で実際に活用されており、タスクの高速化・高精度化に貢献しています。

2022 年 11 月にプレスリリースを出しています。 jp.corp-sansan.com

また、以下の記事でも NineOCR の概要について紹介しています。是非ご覧ください。 buildersbox.corp-sansan.com

NineOCR が抱える課題

リリースを経てプロダクトで活用され始めたことで NineOCR が持つ深層学習モデルの課題がいくつか明らかになりました。

また NineOCR のリリースに向けて研究開発部の体制が強化されたこともあり、数名の研究員で分担しながらモデルの改善に着手できる環境が整ってきました。

メンバーが増え、課題が増えるという状況の変化が起きたものの、開発は EC2 の GPU インスタンス（ml.g4dn.xlarge や ml.g4dn.metal）を数台起動して各メンバーが作業を行う方式を採っていました。

この開発方式は、現状クリティカルではないものの以下の課題が顕在化してきたと感じています。

パラメータチューニング等で一時的に GPU リソースを増やしたいときに、各種コストを加味して躊躇う場面がある
常に GPU リソースを使うとは限らず、必要以上のコストを支払っている場面がある

1 について、必要に応じてインスタンスを増やすことがストレートな解決策ですが、非常に簡素ではあるものの依頼が必要であり作成に若干のリードタイムがあります。

また、インスタンスを作成すると棚卸し作業や OS の更新等で一定のメンテナンスコストがかかります。

中長期的に利用する見立てがあればインスタンスを増やすのが良い選択肢ですが、一時的な用途の場合は作成のコストが上回ると感じる場面がありました。

この課題については、必要なときに必要なだけのリソースを確保でき、インフラストラクチャレベルでのメンテナンスが不要なサービスを利用する、といった解決策が検討できそうです。

2 について、一般に GPU を積んだ EC2 インスタンスの維持コストは比較的高価（例えば g4dn.metal の年間維持コストは約 906 万円*1 ）です。

開発では常に GPU リソースを利用するとは限らず、データ分析や評価、API 仕様の修正等を行う時間もあるため、GPU リソースを使い切れていないと感じる場面がありました。

この課題については、実験計画を綿密に組み立てて空き時間をなくす、使用していない間はインスタンスを停止する等、リソース管理の仕組みを導入する、といった解決策が検討できそうです。

上述の 2 つの課題を解決するサービスとして検討の俎上に上がったのが Amazon SageMaker Training でした。

Amazon SageMaker Training とは

Amazon SageMaker Training の概要は、You Tube の Amazon Web Services Japan 公式チャンネルで公開されている下記動画の 3 分 52 秒あたりに記載されています。

youtu.be

説明を引用します。

「用意したコード」と「用意したデータ」を
「指定したコンピューティングリソース」と「用意した環境」で実行し、
「実行履歴を自動記録」して「アーティファクトを自動で保存する」機能

表現の正確さは失われますが具体化して言い換えてみます。

「手元のコード」と「データ」を
「AWS の GPU インスタンス」と「Docker コンテナ」で実行し、
「ログを S3 に保存」して「モデルも S3 に保存する」機能

端的にはコンテナで処理を実行する機能であり、ECS や Lambda と類似したサービスと捉えて良さそうです。

料金体系も ECS、Lambda と同じく従量課金型であり、処理の実行時間に応じて料金がかかります。したがって、学習を実行している時間が一定以上少なければ EC2 で GPU インスタンスを購入するよりもコスト効率が良くなります。

NineOCR の開発チームにおいては Amazon SageMaker Training がよりコスト効率に優れると試算できたため導入を決定しました。

SageMaker Training の始め方

事前準備

SageMaker Training を実行するためには「学習コード」、「学習データ」、「実行環境」の用意が必要です。

「学習コード」、「学習データ」は通常の EC2 インスタンスで学習を行う場合と同様に準備します。

なお、NineOCR 開発チームでは Feature Store を導入し、学習データの作成を効率化しています。以下の記事でも紹介しておりますので、ぜひご覧ください。

buildersbox.corp-sansan.com

「実行環境」はやや手間ですが、以下のブログポストにある通り大きく分けて 3 つの方法で準備できます。

aws.amazon.com

AWS が提供しているコンテナイメージを拡張する方法
独自のコンテナイメージに SageMaker Training Toolkit をインストールする方法
スクラッチでコンテナイメージを作成する方法

※ 厳密には AWS が提供しているコンテナイメージを"そのまま使う"方法もありますが、 1 のケースに包含されると捉えていただければと思います。

具体的な手順は上記のブログポストに説明があるため割愛し、本エントリーではどの方法を選んだかについて述べます。

なお、選定の観点と評価は筆者の周辺状況を加味した主観的なものであることを予めご了承ください。

選定にあたって重要視したのは、以下の 2 点です。

ローカル（EC2）のホストでも実行できるか
共通処理（ファイルマウント・ログ・エラーハンドリング・S3 への保存等）を実装せずに済むか

1 について補足します。

深層学習のモデル開発は小規模に実験し、コンセプトを確かめるフェーズと大規模に実験し、高精度化を狙うフェーズに大別され、これらを繰り返して試行錯誤するフローと考えています。

小規模に実験する際は、テストやデバッグをしながら高速に実装を進めたいです。

また、一部の設定（学習率スケジューラの挙動・バッチサイズ等）は実際に学習しながら決定したいです。

このフェーズではアジリティを担保するために EC2 等のローカル環境で開発を行うのが望ましいと考えます。

一方で大規模に実験するフェーズでは多少のアジリティは犠牲にして実験を複数同時に実行できることが望ましいと考えます。

試行錯誤を高速に行うために、ローカル環境で小規模に実験する際と大規模に実験する際でコードの変更が極力不要な構成にしたく、この観点を重視しています。

2 については共通処理をマネージドサービスに委ねることで開発のボリュームが減って試行錯誤のフィードバック速度が上がることを期待し、重視しています。

では 1、 2 の観点と 3 つの方法の関係性を確認します。

AWS が提供しているコンテナイメージを拡張する方法

AWS が提供しているイメージは、任意の Python バージョン、深層学習フレームワークのバージョンが提供されているとは限りません*2。

ローカル環境と利用可能なイメージのバージョンに差異がある場合、ローカル環境側のバージョンを変更しなければならないため、ローカルでの実行はハードルが高い場合があります。

一方で、提供されたコンテナイメージは、エラー処理やログ出力に加え複数 GPU での学習やマネージドスポットトレーニング*3にも対応しており、共通処理の実装コストは非常に小さく済みます。

独自のコンテナイメージに SageMaker Training Toolkit をインストールする方法

ローカルの OS と合ったイメージを選択し、バージョン管理ツール経由でライブラリをインストールすることでローカルとコンテナイメージの環境差異をほとんどなくすことができます。

共通処理は SageMaker Training Toolkit が多くを担うため、お作法に従った記述を強いられるものの実装コストを抑えることができます。

スクラッチでコンテナイメージを作成する方法

共通処理をサポートする機能はないので全て独自で実装する必要があります。

まとめ

以下の表に各方法が観点をどの程度充足しているかをまとめます。

観点	AWS が提供しているイメージを拡張	独自のイメージに SageMaker Training Toolkit をインストール	スクラッチでイメージを作成
ローカルのホストでも実行できる	△	◎	◎
共通処理を実装せずに済む	◎	◯	△

上記から、観点の充足度が最も高い独自のコンテナイメージに SageMaker Training Toolkit をインストールする方法を採用しました。

このパターンで環境を用意する場合の Dockerfile は例えば以下の通り作成できます。

バージョン管理ツールには Poetry を利用しています。

FROM nvidia/cuda:11.7.1-base-ubuntu22.04

ENV \
    # os
    DEBIAN_FRONTEND=noninteractive \
    LANG=C.UTF-8 \
    TZ=Asia/Tokyo \
    # python
    PYTHONDONTWRITEBYTECODE=1 \
    PYTHONUNBUFFERED=1 \
    PYTHONUTF8=1 \
    # pip
    PIP_NO_CACHE_DIR=off \
    PIP_DISABLE_PIP_VERSION_CHECK=on \
    # poetry
    POETRY_VERSION=1.3.2 \
    POETRY_HOME=/opt/poetry \
    POETRY_VIRTUALENVS_CREATE=false \
    POETRY_NO_INTERACTION=1 \
    # sagemaker
    SAGEMAKER_PROGRAM=train.py

RUN \
    apt-get update && \
    apt-get install -y --no-install-recommends software-properties-common gcc-x86-64-linux-gnu tzdata curl && \
    add-apt-repository ppa:deadsnakes/ppa -y && \
    apt-get update && \
    apt-get install -y --no-install-recommends python3.10-full python3.10-dev python-is-python3 && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/* && \
    curl -sSL https://install.python-poetry.org | python -

WORKDIR /app

COPY pyproject.toml poetry.lock ./
RUN \
    ${POETRY_HOME}/bin/poetry install --no-cache --no-root --only main,sagemaker-training && \
    rm -rf ~/.cache ${POETRY_HOME}

また、pyproject.toml は例えば以下の通り作成できます。（説明のため、一部記述を省略しています。）

[tool.poetry.dependencies]
python = "^3.10"
pytorch-lightning = {version = "^1.9.0", extras = ["extra"]}

[tool.poetry.group.sagemaker-training]
optional = true

[tool.poetry.group.sagemaker-training.dependencies]
sagemaker-training = "^4.4.5"

学習ジョブの実行

「学習コード」、「学習データ」、「実行環境」が用意できたら、学習ジョブを実行することができます。

なお、以下のコードや説明は独自のコンテナイメージに SageMaker Training Toolkit をインストールする方法で実行環境を用意していることを前提としています。

学習ジョブは boto3 からも実行できますが、Amazon SageMaker Python SDK*4 を利用するとよりシンプルに実行できます。

from datetime import timedelta
from pathlib import Path

from sagemaker.estimator import Estimator
from sagemaker.s3 import S3Uploader


def train(
    image_uri: str, 
    role: str, 
    bucket: str,
    instance_type: str, 
    datasets: Path,
) -> None:

    estimator = Estimator(
        image_uri=image_uri,  # ECR にプッシュされた実行環境の URI
        role=role,  # Training Job の実行ロール
        instance_count=1,
        instance_type=instance_type,
        max_run=int(timedelta(days=5).total_seconds()),
        entry_point="src/train.py",
        dependencies=["src", "config"],
    )

    job_name = estimator._get_or_create_name()

    datasets_s3_uri = S3Uploader.upload(datasets.as_posix(), f"s3://{bucket}/{job_name}/{datasets.stem}")

    estimator.fit(datasets_s3_uri, job_name=job_name)

上記の関数を実行することで src/train.py が SageMaker Training で実行されます。

厳密には、まず src/train.py が SageMaker Training ジョブのコンテナインスタンス上（/opt/ml/code/train.py）にコピーされます。

その後、ジョブは環境変数 SAGEMAKER_PROGRAM に指定されたファイル（上述の Dockerfile の場合 train.py）を、/opt/ml/code から実行します。

したがって Dockerfile で指定した環境変数 SAGEMAKER_PROGRAM と同名のファイルを Estimator の引数である entry_point に指定する必要があります。

また、ローカルのソースコードは Estimator の引数 dependencies に指定することでコンテナへの持ち込みが可能です。

例えば、以下の通り学習に関連するスクリプトを src 配下に格納しておくことで train.py からファーストパーティライブラリを import できるようになります。

src/
├── __init__.py
├── datamodules.py
├── datasets.py
├── models.py
├── schema.py
└── train.py

上記のような構成を取ることで、ローカルでは src/train.py をそのまま実行し、SageMaker Training では上述のスクリプトを実行すれば良く、シームレスに実行環境の切り替えが可能になります。

学習結果の確認

学習のジョブが終了すると以下の通り S3 にモデルやログが出力されます。（一部の情報はマスクしています。）

$ aws s3 ls sagemaker-<region>-<account_id>/<job_name>/ --recursive | grep -v json
9999-12-31 23:59:59          0 <job_name>/debug-output/training_job_end.ts
9999-12-31 23:59:59  190786549 <job_name>/output/model.tar.gz
9999-12-31 23:59:59          0 <job_name>/profiler-output/framework/training_job_end.ts
9999-12-31 23:59:59          0 <job_name>/profiler-output/system/training_job_end.ts
9999-12-31 23:59:59      24341 <job_name>/source/sourcedir.tar.gz
9999-12-31 23:59:59    9591079 <job_name>/training/test.csv
9999-12-31 23:59:59   76720289 <job_name>/training/train.csv
9999-12-31 23:59:59    9589609 <job_name>/training/valid.csv

コンテナインスタンス内の /opt/ml/model 配下に格納されたファイルが圧縮されて <job_name>/output/model.tar.gz として保存されます。

また、ジョブの実行時に指定した依存ファイル類は圧縮されて <job_name>/source/sourcedir.tar.gz として保存されます。

なお、特定のディレクトリ配下に格納されたファイル以外は自動的に S3 に送られることはなく、コンテナの終了とともに消えるため注意が必要です。

詳細は SageMaker Training Toolkit の実装*5を参照していただければと思います。