LightGBMの機械学習モデル予測シミュレーンアプリをFlaskで構築し爆速でHerokuに公開する

本投稿は TECOTEC Advent Calendar 2021 の2日目の記事です。

お世話になっております。テコテックCTOの川人でございます。本ブログには二回目の登場(一回目はこちら)となります。
業務ではめっきりコードを書かなくなりました(= 優秀な皆様にお任せしております)ので、プライベートで開発したものについて記事を書こうと思います。

作ったもの

今回は、機械学習を用いてCOVIDのリスクを予測できるシミュレータを開発しました。実際に公開されており、URLは下記となります。

https://risk-model.herokuapp.com/covid

きっかけは、研究者の知人から論文に引用するためにシミュレータを作りたい、と相談を受けたことで、論文自体も公開されていますので、興味がおありの方はご一読ください。

Prediction of in-hospital mortality with machine learning for COVID-19 patients treated with steroid and remdesivir

前提

本記事にあたっては、以下が揃っていることを前提としています。

また、注意事項は以下の通りです。

  • 開発を行ったのが2021年初となりますので、各種バージョンなどが古いことがあります。適宜読み替えていただければ幸いです。
  • 記事用にソースコードを改変している箇所があります。実際に動作しているものとは異なることをご了承ください。

方針を決める

タイトルにもあるように「爆速で」公開することを目指し、まずは動作するものを作る方針としました。

技術要素としては、機械学習を扱うので言語はPython、フレームワークは軽量であるFlaskを選定(Bottleと悩みましたが実績豊富な技術を選択)しました。

また、無料で運用したいため、PaaSとしてHerokuを用いることにしました。

Flaskアプリを作成する

それでは早速開発していきます。まずはFlaskアプリの雛形を作成します。ファイル構成は下記の通りです。

.
├── Dockerfile
├── app
│   ├── app.py
│   └── requirements.txt
└── docker-compose.yml

最初に、ベースとなるDockerfileは以下のように記述します。

FROM python:3.8.6-slim-buster

COPY ./app /app

WORKDIR /app

RUN pip3 install --upgrade pip \
    && pip3 install -r /app/requirements.txt --no-cache-dir

ENTRYPOINT ["python3"]

CMD ["app.py"]

OSは、Pythonとalpineの相性が良くない、等の情報があった(検証はしておりません)ので、busterを選択(記事執筆時点で、Debian系の最新はbullseyeとなります)。パッケージ管理は、Poetryなどが主流のようですが、Dockerとの共存が大変そうだったため、昔ながらのpipを用いました(スピード重視)。

続いて、app/app.pyを作成します。

import os
from flask import Flask

app = Flask(__name__)
port = int(os.environ.get('PORT', 5000))

@app.route('/')
def index():
    return 'Hello, Flask!'

if __name__ == '__main__':
    app.run(
        debug=bool(os.environ.get('DEBUG', False)),
        host='0.0.0.0',
        port=int(os.environ.get('PORT', 5000)))

debugやportの値を環境変数から取得していますが、debugについては環境によって切り分けたかったため、portについてはHeroku側で動的に生成されるため、そのようにしています。

app/requirements.txtには、Flaskを指定します。

Flask==1.1.2

最後にdocker-compose.ymlを作成します。こちらについては必ずしも必要はありませんのでお好みでご利用ください。

version: '3.8'
services:
  web:
    build: .
    environment:
      DEBUG: 1
      PORT: 5000
    ports:
      - 5000:5000
    volumes:
      - ./app:/app

それではビルド&起動してみましょう。

docker-compose up --build

起動したら、ブラウザから http://localhost:5000 にアクセスし、Hello, Flask! が表示されていればOKです。

Herokuに公開する

それでは、ここまでで作ったアプリをHerokuに公開していきましょう。 下記の公式ドキュメントにある通り、heroku.ymlを作成するだけでDockerイメージをビルドしてくれます(便利な時代です)。

heroku.yml を使用して Docker イメージをビルドする | Heroku Dev Center

前準備として、Herokuの環境を整えましょう。まずHerokuへログインします。

heroku login

上記のコマンドを打つとブラウザが起動するので、作成済みのアカウントでログインしてください。

続いて、Heroku上にアプリケーションを作成します。

heroku create YOUR_APP_NAME

YOUR_APP_NAME には任意の値を設定してください。YOUR_APP_NAME.herokuapp.com というドメインが割り当てられるのでわかりやすい名前が望ましいです。

また、Heroku上のGitにpushすることがデプロイのトリガーとなります(こちらを参照)ので、リモートリポジトリに追加します。

heroku git:remote -a YOUR_APP_NAME

最後に、スタック(詳細はこちら)にコンテナを用いるので明示的に指定しておきます。

heroku stack:set container

これで準備が整いました。heroku.yml を作成していきましょう。

build:
  docker:
    web: Dockerfile

Dockerfileのパスを指定するだけのシンプルな設定となりました。
ここまでの変更をコミットして、Heroku上にPushします。

git push heroku master

これでデプロイが完了しますので、https://YOUR_APP_NAME.herokuapp.com へアクセスしてみてください。先ほどと同様、Hello, Flask! が表示されていると思います(爆速)。

画面を作る

Flask上でフォームを作成するには、WTFormsが良く用いられるようでしたので、今回もそちらを用いることとしました。Flask用のパッケージは Flask-WTF となります。

app/requirements.txt に追記してインストールを行います。

Flask-WTF==0.14.3

パッケージをインストールするには、再ビルドを行うのが確実ですので、現在起動しているコンテナをCtrl+Cで停止させた後、ビルド&起動します。

docker-compose up --build

それではフォームを作成していきます。app/form.py を新たに作ります。

from flask_wtf import FlaskForm
from wtforms.fields import IntegerField, RadioField
from wtforms.validators import AnyOf, InputRequired,  NumberRange, Optional

class COVIDForm(FlaskForm):
    age = IntegerField(
        'Age (years)',
        validators=[
            Optional(),
            NumberRange(min=0)])
    htn = RadioField(
        'HTN',
        coerce=int,
        choices=[(1, 'Yes'), (0, 'No')],
        default=1,
        validators=[
            InputRequired(),
            AnyOf([1, 0])])
    icu = RadioField(
        'ICU',
        coerce=int,
        choices=[(1, 'Yes'), (0, 'No')],
        default=1,
        validators=[
            InputRequired(),
            AnyOf([1, 0])])
    ever_intubated = RadioField(
        'Intubated?',
        coerce=int,
        choices=[(1, 'Yes'), (0, 'No')],
        default=1,
        validators=[
            InputRequired(),
            AnyOf([1, 0])])
    spo2 = IntegerField(
        'Oxygen Saturation (%)',
        validators=[
            Optional(),
            NumberRange(min=0)])
    bun = IntegerField(
        'BUN',
        validators=[
            Optional(),
            NumberRange(min=0)])

基本的には公式ドキュメントのデフォルト通りですが、ラジオボタンの入力値を整数で受け取りたかったので、coerce=intとした(デフォルトでは文字列となります)ところが異なります。

続いてフォームを使えるように app/app.py を修正していきます。主な修正点は3つで、まず、先頭の import を行っている箇所を修正します。

# render_template, redirectを追加します。
from flask import Flask, render_template, redirect

# 下記を追加します。
form import COVIDForm

続いて、CSRFに必要なSECRET_KEYの設定を追記します。

app.config['SECRET_KEY'] = 'YOUR_SECRET_KEY'

YOUR_SECRET_KEYには os.urandom(24) などで生成したものを設定してください。

また、Hello, Flask!を表示していた箇所を下記に変更します。

@app.route('/', methods=['GET', 'POST'])
def index():
    form = COVIDForm()
     if form.validate_on_submit():
        app.logger.debug('success')
    return render_template('index.html', form=form)

これでフォームを利用できるようになりました。最後に app/tempates/index.html を作成します。

<form action="/" method="POST">
  {% for field in form %}
    {% if field.type == 'CSRFTokenField' %}
      {{ field }}
    {% else %}
      <div>
        {{ field.label }}
        {% if field.type == 'RadioField' %}
          {% for subfield in field %}
            {{ subfield }}
            {{ subfield.label }}
          {% endfor %}
        {% else %}
          {{ field }}
        {% endif %}
        {% if field.errors %}
            <ul>
            {% for error in field.errors %}
              <li>{{ error }}</li>
            {% endfor %}
            </ul>
        {% endif %}
      </div>
    {% endif %}
  {% endfor %}
  <input type="submit">
</form>

汎用的になるよう、フィールドによって値を出し分けるようにしています。

ブラウザから再度アクセスしてみると、下記のようなフォームができていると思います。

送信ボタンを押下すると、ログにDEBUG in app: successも出力されていれば正しく動いています(基本的にバリデーションを通過する設定となっています)。

モデルと連動する

いよいよ大詰め、モデルと連携していきます。

LightGBMや機械学習についての説明は割愛とさせていただきますが、知人が学習済みのモデルを渡してくれるので、本アプリでは、そのモデルを読み込み、入力フォームから渡ってきた値をそのモデルに与えて予測値を出し、結果として表示する、という処理を実装していきます。

まず、動作のため、OS上にlightomg.so.1が必要となるので、Dockerfileを修正し、インストールを行います。最終的なDockerfileは下記となります。

FROM python:3.8.6-slim-buster

RUN apt-get update \
    && apt-get upgrade -y \
    && apt-get install -y libgomp1

COPY ./app /app

WORKDIR /app

RUN pip3 install --upgrade pip \
    && pip3 install -r /app/requirements.txt --no-cache-dir

ENTRYPOINT ["python3"]

CMD ["app.py"]

続いて、app/requirements.txtにLightGBMのパッケージを追記してインストールを行います。

lightgbm==3.1.1

次に、入力値から予測した結果を返却するため、app/model.pyを作成します。

import lightgbm as lgb

class COVIDModel:
    def predict(self, X):
        model = lgb.Booster(model_file='covid_model.txt')
        y_pred = model.predict([X])
        return y_pred[0]

X(一次元配列となります)を入力値として受取り、covid_model.txtで定義されたモデルで予測、結果(0-1の値となります)を返却するようなメソッドを持ったクラスを定義しています。

それでは、こちらを用いるように、app/app.pyを修正していきます。こちらについても最終的なものを記します。

import os
from flask import Flask, render_template, redirect

from form import COVIDForm
from model import COVIDModel

app = Flask(__name__)
app.config['SECRET_KEY'] = 'YOUR_SECRET_KEY'

@app.route('/', methods=['GET', 'POST'])
def index():
    form = COVIDForm()
    y_pred = None
    if form.validate_on_submit():
        model = COVIDModel()
        X = [
            form.age.data,
            form.htn.data,
            form.icu.data,
            form.ever_intubated.data,
            form.spo2.data,
            form.bun.data]
        y_pred = round(model.predict(X) * 100, 1)        
    return render_template('index.html', form=form, y_pred=y_pred)

if __name__ == '__main__':
    app.run(
        debug=bool(os.environ.get('DEBUG', False)),
        host='0.0.0.0',
        port=int(os.environ.get('PORT', 5000)))

ソースコードをみていただければわかると思いますが、入力フォームの値からXを生成し、モデルのpredictメソッドを呼び出しています。そして、結果が0-1で返却されるので、%表記にするための処理を行い、テンプレートへ渡しています。

最後に結果を画面に表示するため、app/index.htmlの最終行に下記を付け加えます。

Risk: {{y_pred}} %

こちらで完成です。動作としては公開されているものとほぼ一緒となるので、見た目などはご自由に整えていただければと思います。

本記事は以上となります。それでは皆様良い年の瀬をお過ごしください!

www.tecotec.co.jp