こんにちは！サーバーサイドエンジニアの mitani です。
B/43のサービスの裏側では機械学習を使っている箇所があり、機械学習用のパイプラインとAPIに全てSageMakerを利用しています。今回はなぜSageMakerを採用したのかと、構築時に意識したことについて紹介したいと思います。
インフラ視点でのMLOpsについては @maaaato さんが記事を書いてくれているので、合わせてご覧ください。
Amazon SageMaker StudioでMLOpsを始めました - inSmartBank
機械学習の推論を行う仕組みについて
機械学習の機能をアプリケーションに組み込むためには、大きく分けると3つの仕組みを準備する必要があります。
- 学習用のデータを作成して何かしらのデータストアに保存する仕組み
- 学習用のデータを使って機械学習のモデルを訓練し検証する仕組み
- 訓練済みのモデルを呼び出して推論を行わせる仕組み
定期的にモデルの訓練をやり直して陳腐化を防ぎたいようなケースがなければ、1と2は最初の1回だけ実行して3の推論用の環境を準備するだけでもスモールスタートは可能です。また、頻度が低かったりデータ量が少ない場合は、1と2はローカルのPC内で完結させることもできます。
一方で3の推論を行わせる仕組みは、本番環境のトランザクションの中で安定的に稼働させる必要があるため、できるだけ堅牢な仕組みを構築したいはずです。
推論の仕組みのパターンについてはAWSの資料に詳しく載っているので、ここでは掻い摘んで3つのパターンを紹介します。
Model Embedding
どこかの環境で訓練させた機械学習のモデルをRailsのようなアプリケーションのどこかに配置し、rumaleのようなgemを使ってRailsの中でモデル読み込み&推論を行うような方法です。
推論用に専用のサーバーなどを準備する必要がなく、アプリケーション内部で行われるため通信エラー等の考慮も必要ありません。
一方で、定期的に自動で訓練し直したモデルを配布し直したりする手間が発生したり、モデルの大きさや推論に必要なコンピューターリソースなどを考慮してアプリケーション側のインフラを準備する必要があり、密結合な形となります。
Model Serving (Model as a Service)
推論用のWeb APIアプリケーションを一つ構築する方法です。モデルを読み込んで推論をして結果を返す程度の処理であればそこまで作り込む必要もないため、PythonだとFlaskのフレームワークを使って実装するケースも多いと思います。
モデルの大きさや必要なコンピューターリソースに応じて適切なインフラを準備でき、呼び出し側からはREST APIを意識すれば良いだけなので疎結合に作ることができます。
一方で推論用のAPIを実装し保守する手間があったり、定期的に自動で訓練し直したモデルの再配布の仕組みなどは作る必要があります。
SaaS利用
最後はSageMakerのようなSaaSを利用する方法です。SageMakerには学習用のデータを貯めるための仕組みや、機械学習の訓練 ~ 評価、推論用のAPIのデプロイといった機械学習に必要な機能が十分に揃っており、シームレスに各ステップを繋いでパイプラインを構築することができます。
Model EmbeddingやModel Servingに比べると料金面で差があったり、SageMakerならではの機械学習の記法を学習したりするコストはあります。
SageMakerを選択した理由
B/43ではモデルの訓練 ~ 評価、推論APIのデプロイまでをSageMakerで構築しています。SageMakerを選択した理由としては次のようなものがあります。
PoC的な側面が強いので導入しやすく剥がしやすい疎結合な形を意識
機械学習の使うサービス自体が、今後もずっと機械学習を組み込み続けたいという訳ではなく、一旦効果があるのかの検証をしてみたいというフェーズでした。そのため、できるだけ疎結合な形で実装して、効果がなくてやめたい時に剥がしやすい構成を意識しました。
Model Embeddingの方式だと、同じリポジトリを触る他のエンジニアに理解したもらったり意識してもらう必要があり、機械学習用のgemを入れたりと密結合になってしまう点が気になりました。またModel Serving方式では専用の推論APIを実装するのが手間で、そこに工数をかけたくありませんでした。
そのため、最初の知識習得のコストはかかってしまいますがSageMakerを利用して、疎結合な構成とすることにしました。
訓練 ~ デプロイまでのパイプラインの構築が簡単
SageMakerではJupyter Notebookでモデルの訓練から評価、デプロイまでの処理をコードで書くだけでパイプラインを構築することができます。
訓練するためのコードの書き方はSageMaker独自の書き方になっているため、公式のドキュメントに書かれていないようなことをやりたい時には、多少調査に時間がかかったりはします。ただ、公式のexamplesがかなり豊富にあり、よほど独自の前処理を挟んだりするケースがなければほとんどexamples通りに実装するだけでデプロイまで持っていけると思います。
今回機械学習を導入したいサービスではある程度データの鮮度が重要だと考えていたため、訓練~デプロイまでを一気通貫で一つのSaaS内で完結できるのは魅力的でした。
SageMakerで構築した各パイプライン
ここからは具体的に構築した内容を紹介しようと思います。
学習データの準備
学習データの準備はSageMakerの外側でやっています。Railsのrakeタスクで特徴量のデータを集計してcsvとして出力し、訓練用のデータをおくS3バケットにアップロードするタスクを準備しています。
SageMakerはほんとに多機能が備わっておりAmazon SageMaker Feature StoreやAmazon SageMaker Data Wranglerといった機能を使って特徴量を管理することもできます。ただ多機能すぎて使いこなすために時間がかかってしまい、早期に立ち上げてリリースしたいPJではToo Much感があったため利用を見送っています。
学習
S3にファイルがアップロードされるとSageMakerのPipelineが起動されるようになっています。このPipelineの中ではS3からファイルをダウンロードして訓練し、モデルの評価を行って一定の品質以上であればSageMakerのModelRegistoryにモデルを保存するようにしています。
モデルが保存されたら、エンジニアが目視でモデルの性能を評価するようにしています。Pipelineでは例えばAOCの値が一定より悪ければPipelineを失敗させるような仕組みもありますが、より細かくモデルの性能を検証した上でデプロイするかの判断を行いたかったため、モデルの保存とデプロイのPipelineは分けています。
デプロイ
モデルの性能が十分あると判断できたらModel Registoryでモデルのステータスをapprovedに更新します。approvedに更新されるとEventBridgeでデプロイ用のPipelineが起動されるようになっています。
推論用のAPIにはSageMaker Endopointを利用しています。SageMaker Endpointはデプロイ方法をBlue/Greenやローリングから選ぶことができ、ここら辺の機能が揃っているのも本番運用がやりやすいポイントです。B/43ではBlue/Greenデプロイを選択しています。
構築時に気をつけたこと
現在の機械学習モデルは効果検証に重きを置いていることと、構築した私自身がサーバーサイドエンジニアで機械学習専任ではないため、できるだけシンプルなアーキテクチャとして、私以外のエンジニアが担当することになってもSageMakerのドキュメントを見ればどの技術が使われているのかをそこまで難しくなく理解できるようにしました。
また、訓練したモデル自体もあまり特徴量にこだわり過ぎたり、チューニングをやり過ぎたりしないようにしています。今はまだ機械学習の専任者がいないため、あまり尖ったモデルになり過ぎて推論の性能が安定しなかったり、環境変化に左右され過ぎると保守に手間がかかるため、カリカリに精度が出るようなモデルではなく安定して精度にばらつきが出ないようにできるだけ意識して作りました。
終わりに
今回はB/43で使っている機械学習のパイプラインについて紹介しました。SageMakerはほんとに多機能で、今はその一部しか使えていないので、今後は色々な機能を使っていきたいと思います。
また、 @maaaato さんのブログの方にも書いてありますが、SageMakerのMLOpsの情報はある程度ネットで出回っていたり、公式のドキュメントもあるので初期の学習コストは多少高いものの、機械学習をほぼ知らない人でも環境構築ができるので、完全に0から作ることに比べるとかなりの工数削減ができたかと思います。
とはいえまだまだ改善の余地しかないので、一緒に作っていける人を絶賛募集中です！
