🔧

SageMaker で学習 - 推論間の前処理を共通化する

2022/07/13に公開

AWS

MLOps

SageMaker

tech

こんにちは。エンジニアチームの山岸 (@yamagishihrd) です。

SageMaker を使用した機械学習（以下、「ML」）における前処理の実装方式について検討してみたので、今回はその内容について紹介したいと思います。

1. ML における前処理

a. バッチ前処理（学習データ作成時）

通常、ML で利用する学習データには前処理が必要になります。例えば、自然言語処理のユースケースの場合、少なくとも単語列からなる文章に対して「単語埋め込み (Word embedding)」を適用し、数値配列として扱えるようにする必要があるでしょう。

他にも、ML モデルの学習に有効な特徴量を作成するための前処理が必要になるかもしれません。（例: 「文章が特定の単語を含んでいるか否かを示すフラグを付与する」）

これら学習用データセット作成のための前処理を、まとめて「バッチ前処理」と呼ぶことにします。SageMaker Processing という機能を利用すると、前処理・後処理・評価といった任意のワークロードを SageMaker マネージドな環境で実行することができます。

実装例

preprocessing / Dockerfile

FROM python:3.9

RUN pip install -U pip
RUN pip install numpy pandas torch lightgbm 
RUN pip install transformers pytorch-lightning
RUN pip install ipadic fugashi

COPY preprocess.py /opt/ml/code/

ENV PYTHONUNBUFFERED=TRUE
ENTRYPOINT ["python3", "/opt/ml/code/preprocess.py"]

preprocessing / preprocess.py

# coding: utf-8
import sys, os
import json

import numpy as np
import pandas as pd
import torch
from transformers import BertJapaneseTokenizer, BertModel

prefix = "/opt/ml"
input_prefix = os.path.join(prefix, "processing/input")
output_prefix = os.path.join(prefix, "processing/output")

MODEL_NAME = "cl-tohoku/bert-base-japanese-whole-word-masking"

class BertEncoder:
    
    def __init__(self):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.max_len = 256
        self.tokenizer = BertJapaneseTokenizer.from_pretrained(MODEL_NAME)
        self.bert_model = BertModel.from_pretrained(MODEL_NAME).to(self.device)

    def vectorize(self, document: str) -> np.array:
        # 略
        return vector

def transform(inp, encoder=None):
    encoded = encoder.vectorize(inp["document"]).tolist()
    features = add_some_features(inp["document"]) # 任意の処理
    transformed = encoded + features
    return transformed

if __name__ == "__main__":

    filepath = os.path.join(input_prefix, "path/to/input_file.csv")
    df = pd.read_csv(filepath)

    encoder = BertEncoder()
    records = list()
    for index, row in df.iterrows():
        record = [row["label"]]
        record.extend(transform(inp=row["document"], encoder=encoder))
        records.append(record)

    records_df = pd.DataFrame(np.array(records))
    records_df[0] = result_df[0].apply(lambda x: int(x)) # ラベル列を int 型にキャスト

    output_filepath = os.path.join(output_prefix, "path/to/output_file.csv")
    result_df.to_csv(output_filepath, header=False, index=False)

b. リアルタイム前処理（ML モデル呼び出し時）

学習済み ML モデルを、リアルタイム推論 API として機能させるケースを考えます。このとき、推論時にもバッチ前処理と同様の前処理を、入力に対して適用する必要があります。リアルタイム推論における前処理を、「リアルタイム前処理」と呼ぶことにします。

2. 課題感

前処理に対する課題感としては「バッチ前処理とリアルタイム前処理をいかに共通化させるか」というものです。バッチ前処理のロジックを推論用スクリプトにも実装してしまうと、同じコードを二重で管理することになり、運用していく中で両者に乖離が生じてしまうかもしれません。

同じコードをバッチ前処理／リアルタイム前処理の双方から共通的に利用するには、どう実装するのが良いでしょうか。

3. 実装

上記課題への対策として、「SageMaker Processing で利用するバッチ前処理用スクリプトを、推論用コンテナにも同居させて呼び出す」構成としました。詳細について、以下に解説します。

リポジトリの構成

学習・推論用コンテナの実装 | Dockerfile

学習・推論用コンテナは、「独自コンテナ (Bring-Your-Own Container)」^[1] で実装します。（本記事のベースとなる公式サンプルが GitHub 上に公開されています ^[2]）

Dockerfile の実装例は以下の通りです。ポイントは、前処理用コンテナのエントリポイントである preprocessing/preprocess.py を、学習・推論用コンテナにも COPY している点です。

Dockerfile

FROM python:3.9

RUN apt-get -y update
RUN apt-get -y install nginx ca-certificates
RUN pip install -U pip
RUN pip install flask gunicorn gevent
RUN pip install numpy pandas schikit-learn torch lightgbm
RUN pip install ipadic fugashi transformers

COPY model /opt/ml/code
COPY preprocessing/preprocess.py /opt/ml/code

RUN chmod +x /opt/ml/train
RUN chmod +x /opt/ml/serve
ENV PYTHONBUFFERED=TRUE
ENV PYTHONDONTWEITEBYTECODE=TRUE
ENV PATH="/opt/ml/code:${PATH}"
WORKDIR /opt/ml/code

リアルタイム前処理の実装 | predictor.py

推論処理を実装している model/predictor.py の中で、preprocess.transform を import します。バッチ前処理 (SageMaker Processing ジョブ) では main 関数から呼び出していましたが、関数として切り出しておくことでリアルタイム前処理からも利用できます。

predictor.py

# coding: utf-8
# 〜略〜
import pickle
import flask
from preprocess import transform, BertEncoder

encoder = BertEncoder()

class ScoringService(object):
    model = None
    @classmethod
    def get_model(cls):
        if cls.model == None:
            with open("/opt/ml/model/model.pkl"), "rb") as f_model:
                cls.model = pickle.load(f_model))
        return cls.model
    @classmethod
    def predict(cls, inp):
        clf = cls.get_model()
        return clf.predict(transform(inp))
# 〜略〜

まとめ

「バッチ前処理（学習時）とリアルタイム前処理（推論時）のコードをいかに共通化するか」という課題に対し、実装方式を検討してみました。結論としては、「バッチ前処理スクリプトを推論用コンテナにも同居させて呼び出す」という構成で実現できました。

以上、SageMaker 開発 Tips でした！

脚注