Amazon SageMakerを利用した効率的な機械学習 with Rust

asada

作成: 2021.11.02 更新: 2021.11.02

はじめに
SageMakerでの4つの学習/推論処理パターン
Rust With SmartCoreでの学習処理
Rust With actix-webでの推論処理
Dockerイメージを作成する
SageMakerから、学習処理、推論処理を実行する
SageMakerをローカルモードで実行する
終わりに

はじめに

デジタルイノベーション部の浅田です。

クラウドを利用した開発を行うにあたって、クラウドを上手く利用しようとすればするほど、ローカル開発環境と本番環境（クラウド環境）とでの実装方法の差分を少なくすることが効率的に開発を行う上で重要になってきます。

例えば、Amazon DynamoDBを利用してサービスを開発しようとすると、ローカル開発環境でどのように開発を進めるか？という課題が生まれます。DynamoDBであれば、ローカルのエミュレータが提供されているので、それを利用するという解決策が考えられます。

機械学習においても、ローカル開発環境と本番環境とのやり方を統一できたほうが、効率的に開発ができます。

その一つのやり方が、Amazon SageMaker（以下SageMaker）を利用することで、ローカル環境と本番環境とで差分の少ない、統一的な方法で開発することです。

また、機械学習においては学習処理と推論処理とが存在し、学習処理で作成した機械学習モデルを推論処理で利用するといったことが行われますが、SageMakerを利用することで、言語やフレームワークに依存しない形でそれぞれの処理を連携させることができます。

そこで、今回はRustでの機械学習をSageMakerで行う例を通して、

ローカル環境と本番環境とで統一的な方法で開発できる
言語やフレームワークに依存しない方法で学習処理や推論処理の連携ができる

ということをご紹介したいと思います。

SageMakerでの4つの学習/推論処理パターン

SageMakerで機械学習モデルを学習するにあたって、大きく分けると4パターンあります。

Auto Pilotを使って、完全にSageMakerに任せる
組み込みアルゴリズムやJumpStartを利用して用意されているアルゴリズムを選んで利用
組み込みのフレームワーク（Tensorflow, Pytorch, XGBoost, etc.）を利用
実行環境やコードをDockerイメージで用意

今回は4つ目のパターンになります。実行環境をDockerイメージにすることで、SageMakerに用意されていないような言語、フレームワークを利用して学習、および推論を行うことができます。

今回はRustでSmartCoreというフレームワークを利用して学習、および推論を行います。そして、それはRust & SmartCoreに限った方法ではないので、他の言語、フレームワークにも応用が利く方法になります。

Rust With SmartCoreでの学習処理

SmartCoreはRustで実装された機械学習フレームワークになります。Pythonの機械学習フレームワークであるscikit-learnと似たライブラリになっているので、scikit-learnを使いなれた方には、なじみやすいフレームワークとなっています。

以下が、SmartCoreを利用した学習処理のコードです。

use ndarray::prelude::*;
use ndarray::{Array, OwnedRepr};
use ndarray_csv::Array2Reader;
use smartcore::linear::logistic_regression::LogisticRegression;
use smartcore::metrics::accuracy;
use smartcore::model_selection::train_test_split;
use std::error::Error;
use std::fs::File;
use std::io::prelude::*;

const DATA_PATH: &str = "/opt/ml/input/data/training/iris.data";
const MODEL_PATH: &str = "/opt/ml/model/lr.model";

// CSVの読み込み
fn read_csv_to_array2() -> Array2<String> {
    let mut rdr = csv::ReaderBuilder::new()
        .has_headers(false)
        .from_path(DATA_PATH).expect("Can not read csv.");
    rdr.deserialize_array2_dynamic().unwrap()
}

// 特徴量のデータを取得
fn get_features(data :&Array2<String>) -> ArrayBase<OwnedRepr<f32>, Dim<[usize; 2]>> {
    let input = data.slice(s![.., 0..4]);
    let mut vec_x: Vec<f32> = Vec::new();
    for i in input.iter() {
        vec_x.push(i.parse().unwrap());
    }
    Array::from_shape_vec( (data.nrows(), 4), vec_x ).unwrap()
}

// 正解ラベルのデータを取得
fn get_target(data :&Array2<String>) -> ArrayBase<OwnedRepr<f32>, Dim<[usize; 1]> > {
    let target = data.slice(s![.., 4]);
    let mut vec_y: Vec<f32> = Vec::new();
    for t in target.iter() {
        let t_f = match t.as_str() {
            "Iris-setosa" => 0.,
            "Iris-versicolor" => 1.,
            "Iris-virginica" => 2.,
            _ => 0.,
        };
        vec_y.push(t_f);
    }
    Array::from_shape_vec(data.nrows(), vec_y).unwrap()
}

// 学習済みモデルを保存
fn save_model(model: &LogisticRegression<f32, ArrayBase<OwnedRepr<f32>, Dim<[usize; 2]>>>) -> Result<(), Box<dyn Error>> {
    let model_bytes = bincode::serialize(&model).expect("Can not serialize the model");
    File::create(MODEL_PATH)
        .and_then(|mut f| f.write_all(&model_bytes))
        .expect("Can not persist model");
    Ok( () )
}

// メイン処理
fn main() -> Result<(), Box<dyn Error>> {
    let data = read_csv_to_array2();
    let (x, y) = (get_features(&data), get_target(&data));
    let (x_train, x_test, y_train, y_test) = train_test_split(&x, &y, 0.3, true);
    let model = LogisticRegression::fit(&x_train, &y_train, Default::default()).unwrap();
    let y_hat = model.predict(&x_test).unwrap();
    println!("accuracy: {}", accuracy(&y_test, &y_hat));
    save_model(&model)
}

ここでは、ロジスティック回帰による分類タスクを実装しています。学習データは、UCI Machine Learning Repository: Iris Data Setを利用させていただきました。

SageMakerで学習処理を行うにあたって重要な点は2点です。

入力データは/opt/ml/input/data/training/配下にSageMakerによって格納されます。後に示すように、SageMakerではローカルのファイルや、S3上のファイルを学習データとして指定することができますが、SageMakerが学習処理コンテナ上の/opt/ml/input/data/training/に配置してくれるので、学習処理側ではファイルがどう指定されるかを気にする必要はありません。
学習したモデルを/opt/ml/model/配下に保存します。学習処理はコンテナ上で行われるので、学習したモデルは推論処理のためにコンテナ外に退避する必要があります。それはローカルであったり、S3であったりするのですが、学習処理側としては/opt/ml/model/配下に置いておけば、SageMakerがあとは保存の面倒を見てくれます。

Rust With actix-webでの推論処理

次は推論処理になります。SageMakerで推論処理のエンドポイントを立ち上げるにあたって、いくつかの条件を満たす必要がありますが、その一つにHTTPアクセスのインターフェースを実装する必要があります。早い話がWebアプリケーションを用意する必要があるということです。

なので、今回はactix-webというRustのWebアプリケーションフレームワークを利用します。

以下が、actix-webを利用した推論処理のコードになります。

use actix_web::{guard, web, App, HttpResponse, HttpServer, Responder};
use lazy_static::lazy_static;
use ndarray::prelude::*;
use ndarray::{Array, OwnedRepr};
use serde::Serialize;
use smartcore::linear::logistic_regression::LogisticRegression;
use std::fs::File;
use std::io::prelude::*;
use std::str;

const MODEL_PATH: &str = "/opt/ml/model/lr.model";

// レスポンスJSON用構造体
#[derive(Serialize)]
struct PredictResult {
    predicted: i32,
}

//　学習済みモデルのロード
lazy_static! {
    static ref MODEL: LogisticRegression<f32, ArrayBase<OwnedRepr<f32>, Dim<[usize; 2]>>> = {
        let mut buf: Vec<u8> = Vec::new();
        File::open(MODEL_PATH)
            .and_then(|mut f| f.read_to_end(&mut buf))
            .expect("Can not load model");
        bincode::deserialize(&buf).expect("Can not deserialize the model")
    };
}

// ヘルスチェック用
async fn ping() -> impl Responder {
    HttpResponse::Ok()
}

// 推論処理用
async fn invocations(body: web::Bytes) -> impl Responder {
    let csv = str::from_utf8(&body).unwrap();
    let v: Vec<f32> = csv
        .split(',')
        .map(|s| s.parse().expect("not floating number!"))
        .collect();
    let x = Array::from_shape_vec( (1, v.len() ), v).unwrap();
    let y_hat = MODEL.predict(&x).unwrap();
    HttpResponse::Ok().json(PredictResult {
        predicted: y_hat[0] as i32,
    })
}

// サーバ起動
#[actix_web::main]
async fn main() -> std::io::Result<()> {
    HttpServer::new(|| {
        App::new()
            .route("/ping", web::get().to(ping))
            .route(
                "/invocations",
                web::post()
                    .guard(guard::Header("content-type", "text/csv"))
                    .to(invocations),
        )
    })
    .bind( ("0.0.0.0", 8080) )?
    .run()
    .await
}

Webアプリケーション化するといっても、複雑なことをする必要はありません。

SageMakerで推論処理を実装するにあたってポイントは2点です。

/pingへのgetリクエストに正常に応答すること。
/invocationsへのpostリクエストに推論結果を返すこと。

なお、SageMakerが/opt/ml/model/に学習処理で保存されたモデルを配置してくれるので、推論処理側は外部にあるモデルを取得してくる処理を実装する必要はなく、/opt/ml/model/に保存されたモデルを読み込めばOKです。

Dockerイメージを作成する

さて、学習処理と推論処理のコードは用意できたので、次はそれを利用するためのDockerイメージを作成します。

SageMakerは学習処理の際に、trainというコマンドで実行します。つまり、コンテナ内でtrainと打った時に、前述の学習処理が実行されればよいということになります。同様に推論処理はserveというコマンドを実行します。

Rustの場合、話はかなりシンプルになります。学習処理のコードをtrainというバイナリに、推論処理のコードをserveというバイナリにコンパイルしたうえで、PATH環境変数が通っている場所(例えば/usr/bin)に配置すればよいということになります。

|-- Cargo.lock
|-- Cargo.toml
|-- Dockerfile
`-- src
    `-- bin
        |-- serve
        |   `-- main.rs
        `-- train
            `-- main.rs

上記のような形で、src/bin/trainに学習処理のコードを、src/bin/serveに推論処理のコードを配置した上で、

cargo build --release

とコンパイルを行えば、target/release/train, target/release/serveにバイナリが得られるので、それをDockerイメージの/usr/binに配置します。

上記を踏まえたDockerfileが以下になります。

FROM rust:1.56.0 AS builder
RUN mkdir -p /app
WORKDIR /app
COPY . /app
RUN cargo build --release

FROM debian:bullseye-slim
COPY --from=builder /app/target/release/train /usr/bin/train
COPY --from=builder /app/target/release/serve /usr/bin/serve
EXPOSE 8080

これをECRにpushします。

aws ecr get-login-password --region ap-northeast-1 | docker login --username AWS --password-stdin ${ACCOUNT_ID}.dkr.ecr.ap-northeast-1.amazonaws.com
docker build -t rust-ml:sagemaker .
docker tag rust-ml:sagemaker ${ACCOUNT_ID}.dkr.ecr.ap-northeast-1.amazonaws.com/rust-ml:sagemaker
docker push ${ACCOUNT_ID}.dkr.ecr.ap-northeast-1.amazonaws.com/rust-ml:sagemaker

SageMakerから、学習処理、推論処理を実行する

SageMakerの操作はsagemakerというpythonライブラリから行いますので、pipでインストールします。その際に、のちのちローカルモードも実行するので、sagemaker[local]としてインストールします。

python3 -m venv exec-env
. exec-env/bin/activate
pip install sagemaker

あらかじめ、S3バケット(今回の場合であれば、sagemaker-with-rust)にUCI Machine Learning Repository: Iris Data Setよりダウンロードしたiris.dataを配置しておきます。

SageMakerでの学習、および推論処理を行うコードは以下になります。

import sagemaker
from sagemaker.estimator import Estimator

session = sagemaker.Session()
account_id = session.boto_session.client('sts').get_caller_identity()['Account']
role = 'arn:aws:iam::{}:role/SageMakerExecutionRole'.format(account_id)

# 学習データの配置場所
training = 's3://sagemaker-with-rust/training'
# 学習済みモデルの配置場所
output = 's3://sagemaker-with-rust/output'

# 推論器の作成
est = Estimator(
    image_uri=account_id+'.dkr.ecr.ap-northeast-1.amazonaws.com/rust-ml:sagemaker',
    role=role,
    instance_count=1,
    instance_type='ml.m5.large',
    output_path=output,
)

# 学習処理
est.fit({'training':training})

# 推論エンドポイントのデプロイ
pred = est.deploy(instance_type='ml.t2.medium', initial_instance_count=1)

# サンプルデータで推論を呼び出し
pred.serializer = sagemaker.serializers.CSVSerializer()
pred.deserializer = sagemaker.deserializers.JSONDeserializer()
test_samples = ['7.2,3.0,5.8,1.6']
response = pred.predict(test_samples)
print(response)

# 終わったら削除
pred.delete_endpoint()

ちなみに各自でコードを実行する際にはS3のバケット名、およびroleに指定しているSageMaker実行用ロールは、ご自身のアカウントのものに変更する必要があります。

ポイントとしては、以下になります。

Estimatorの作成時にimage_urlとしてpushしたECR上のイメージを指定
Estimatorの作成時にoutputとして学習済みのモデルを配置
- 学習処理で、/opt/ml/model/配下に保存したモデルがSageMakerによって指定した場所に保存されます
trainingデータの配置場所をfitメソッドのコール時に指定
- SageMakerによって、学習処理コンテナ上の/opt/ml/input/data/training/に学習データのファイルが配置されます

上記のコードをsagemaker-remote.pyとして保存し、

python sagemaker-remote.py

を実行して、以下のような結果が出力されれば成功です。

2021-10-24 06:19:12 Starting - Starting the training job...
2021-10-24 06:19:35 Starting - Launching requested ML instancesProfilerReport-1635142752: InProgress
......
2021-10-24 06:20:35 Starting - Preparing the instances for training......
2021-10-24 06:21:38 Downloading - Downloading input data...
2021-10-24 06:22:12 Training - Training image download completed. Training in progress.
2021-10-24 06:22:12 Uploading - Uploading generated training model
2021-10-24 06:22:12 Completed - Training job completed
.accuracy: 0.9777778
Training seconds: 34
Billable seconds: 34
-----!{'predicted': 2}

サンプルデータとして渡している’7.2,3.0,5.8,1.6’は、Iris-virginicaのデータなので、推論結果もあたっています。

SageMakerをローカルモードで実行する

さて、めでたくSageMakerでRustの機械学習処理を実行できたわけですが、実際に学習処理を行ってから、推論結果を得るまで10分ほどかかっています。また、課金時間としては数十秒ほどかかっています。

兎角、機械学習は試行錯誤の連続なので、このリードタイムはかなりクリティカルですし、課金金額も数十秒とはいえかかっているので、積み重なっていけばそこそこの金額になってしまうのも、避けたいところです。

そこで、SageMakerのローカルモードで実行することで、上記の問題を回避します。

SageMakerのローカルモードは、AWS環境のインスタンスを使う代わりに、自分の端末内でDockerコンテナを使います。実行結果を得るまでの時間も速くなりますし、課金も発生しないので、アルゴリズムを試行錯誤で試す際に向いています。

ローカルモード実行のコードは以下のようになります。

import sagemaker
from sagemaker.estimator import Estimator

session = sagemaker.Session()
account_id = session.boto_session.client('sts').get_caller_identity()['Account']
role = 'arn:aws:iam::{}:role/SageMakerExecutionRole'.format(account_id)

# 学習データの配置場所
training = 'file://.'
# 学習済みモデルの配置場所
output = 'file://.'

# 推論器の作成
est = Estimator(
    image_uri='rust-ml:sagemaker',
    role=role,
    instance_count=1,
    instance_type='local',
    output_path=output,
)

# 学習処理
est.fit({'training':training})

# 推論エンドポイントのデプロイ
pred = est.deploy(instance_type='local', initial_instance_count=1)

# サンプルデータで推論を呼び出し
pred.serializer = sagemaker.serializers.CSVSerializer()
pred.deserializer = sagemaker.deserializers.JSONDeserializer()
test_samples = ['7.2,3.0,5.8,1.6']
response = pred.predict(test_samples)
print(response)

# 終わったら削除
pred.delete_endpoint()

ちなみに、sagemaker-remote.pyと差分がないようにしているので書いてありますが、実際にはroleはダミーでも問題ないので、session ~ roleを取得している3行(4~6行目)は必要ありません。つまり、16行目をrole=”dummy/dummy”とすることで、 4~6行目はなくても問題はありません。

ポイントとしては、3点です。

学習データや学習済みモデルの配置場所をS3ではなく、”file://.”とすることで、ローカルに保存されるようにしている
Estimatorの作成時、image_urlにローカルのイメージ名（ここでは”rust-ml:sagemaker”）を指定する
Estimatorの作成時、およびデプロイする際のinstance_typeに’local’を指定する

SageMakerをローカルモードで動かす時に必要な変更はこれだけです。なので、環境変数などでローカル環境実行時と本番環境実行時との動作を簡単に切り替えることができます。

ローカルモードでの実行のリードタイムは30秒もないので、試行錯誤を手軽に繰り返すことができますし、料金もかかりません。それでいて実際に本番環境に適用するための変更はわずかで済みます。

上記のコードをsagemaker-local.pyという名前で保存し、

python sagemaker-local.py

を実行して、以下のような出力が出れば成功です。

Creating rxv8w31h0l-algo-1-ira5t ...
Creating rxv8w31h0l-algo-1-ira5t ... done
Attaching to rxv8w31h0l-algo-1-ira5t
rxv8w31h0l-algo-1-ira5t | accuracy: 0.93333334
rxv8w31h0l-algo-1-ira5t exited with code 0
Aborting on container exit...
===== Job Complete =====
!{'predicted': 2}
Attaching to 6d8lwisy9v-algo-1-cr4kt
Gracefully stopping... (press Ctrl+C again to force)

なお、accuracyは訓練データとバリデーションデータをランダムに選択しているので、実行の度に変動します。ローカルだから下がっているわけではありません。

終わりに

Rustでの機械学習をSageMakerで行う例を通して、SageMakerでローカルと本番時とでの差分を最小限にしつつ機械学習処理の開発をする方法をご紹介しました。

この方法は言語やフレームワークに縛られないやり方なので、他の様々な言語やフレームワークでも同じ方法で運用することができます。

このようにSageMakerを使うことで、ローカル環境と本番環境との差分をなるべく小さくしつつ、様々な言語、フレームワークに対応した統一的な開発、運用を行うことができます。

次回、文章からの固有表現抽出をSageMakerとComprehendを利用して効率的に行う方法をお届けする予定です！