Elixir Phoenix で Bumblebee による画像分類 REST API を構築する

はじめに

Elixir で AI を簡単に動かす Bumblebee を一通り使ってみました

• 画像分類: ResNet50
• 画像生成: Stable Diffusion
• 文章の穴埋め: BERT
• 文章の判別: BERTweet
• 文章の生成: GPT2
• 質疑応答: RoBERTa
• 固有名詞の抽出: bert-base-NER

せっかくなので、これを Phoenix で動かしてみましょう

公式のサンプルはこちら

私は mix phx.new から REST API で構築しました

最終的には AWS の SageMaker 上で画像AIのマイクロサービスとして動かします

私の実装はこちら

実装環境

• macOS 13.1
• Elixir 1.14.2 OTP 25

プロジェクトの作成

以下のコマンドで Phoenix プロジェクトを作成します

LiveView も DB も何も使わない REST API なので全て拒否しています

mix phx.new api \
  --no-assets \
  --no-ecto \
  --no-html \
  --no-gettext \
  --no-dashboard \
  --no-live \
  --no-mailer

api ディレクトリーの配下に Phoenix のテンプレートが生成されます

┬config
├deps
├lib
├test
├mix.exs
├mix.lock
└README.md

不要物の削除

README.md は要らないので削除します

lib/api_web/endpoint.ex から不要な行を削除します

...
- # socket "/live", Phoenix.LiveView.Socket, websocket: [connect_info: [session: @session_options]]
-
- # Serve at "/" the static files from "priv/static" directory.
- #
- # You should set gzip to true if you are running phx.digest
- # when deploying your static files in production.
- plug Plug.Static,
-   at: "/",
-   from: :api,
-   gzip: false,
-   only: ~w(assets fonts images favicon.ico robots.txt)
-
...

REST API として動かすため、静的コンテンツ配信は不要です

依存モジュールの追加

mix.exs に以下の様に依存モジュールを追加します

...
  defp deps do
    [
      {:phoenix, "~> 1.6.15"},
      {:telemetry_metrics, "~> 0.6"},
      {:telemetry_poller, "~> 1.0"},
      {:jason, "~> 1.2"},
+     {:bumblebee, "~> 0.1"},
+     {:stb_image, "~> 0.6"},
+     {:exla, "~> 0.4"},
      {:plug_cowboy, "~> 2.5"}
    ]
  end
...

設定の変更

config.exs に以下の行を加えます

EXLA を Nx のバックエンドとして動かす（処理を高速化する）ためです

...
config :phoenix, :json_library, Jason

+config :nx, default_backend: EXLA.Backend
...

dev.exs を以下の様に変更します

外部からのアクセスを許可し、ポート番号を 8080 にするためです

8080 は SageMaker で外部向けに出すポート番号です

...
config :api, ApiWeb.Endpoint,
  # Binding to loopback ipv4 address prevents access from other machines.
  # Change to `ip: {0, 0, 0, 0}` to allow access from other machines.
- http: [ip: {127, 0, 0, 1}, port: 4000],
+ http: [ip: {0, 0, 0, 0}, port: 8080],
  check_origin: false,
...

コードの追加

serving.ex

lib/api_web/serving.ex を以下の内容で追加します

ここで Bumblebee に提供させたい Nx.Serving を準備しておきます

defmodule ApiWeb.Serving do
  use Agent

  # ResNet50 のリポジトリーIDを指定
  @resnet_id "microsoft/resnet-50"

  def start_link(_opts) do
    # モデルのダウンロード
    {:ok, model} =
      Bumblebee.load_model({:hf, @resnet_id})

    {:ok, featurizer} =
      Bumblebee.load_featurizer({:hf, @resnet_id})

    # ResNet の画像識別サービスを生成
    resnet = Bumblebee.Vision.image_classification(model, featurizer)

    # Agent に入れておく
    Agent.start_link(fn ->
      %{
        resnet: resnet,
      }
    end, name: __MODULE__)
  end

  # 使用時に Agent から取り出す
  def get(key) do
    Agent.get(__MODULE__, &Map.get(&1, key))
  end
end

ローカルで動作させる場合は Bumblebee.load_model に cache_dir を指定しません

次の工程でコンテナ化するときには指定します

application.ex

lib/api/application.ex に以下の内容を追加します

  def start(_type, _args) do
    children = [
+     ApiWeb.Serving,
      # Start the Telemetry supervisor
      ApiWeb.Telemetry,
      # Start the PubSub system
      {Phoenix.PubSub, name: Api.PubSub},
      # Start the Endpoint (http/https)
      ApiWeb.Endpoint
      # Start a worker by calling: Api.Worker.start_link(arg)
      # {Api.Worker, arg}
    ]
...

ping_controller.ex

lib/api_web/controllers/ping_controller.ex を以下の内容で追加します

SageMaker でヘルスチェック（死活確認）をするのに使う Ping 応答用のコントローラーです

defmodule ApiWeb.PingController do
  use ApiWeb, :controller

  action_fallback ApiWeb.FallbackController

  def ping(conn, _params) do
    conn
    |> put_status(200)
    |> json(%{})
  end
end

prediction_controller.ex

lib/api_web/controllers/prediction_controller.ex を以下の内容で追加します

推論 API 用のコントローラーです

先程作った ApiWeb.Serving で定義した画像識別サービスに画像をテンソル化して渡し、推論結果を返します

defmodule ApiWeb.PredictionController do
  use ApiWeb, :controller

  alias ApiWeb.Serving

  action_fallback ApiWeb.FallbackController

  # BASE64 文字列で受ける場合
  def index(conn, %{"image" => base64}) do
    predictions =
      base64
      |> Base.decode64!()
      |> predict()

    conn
    |> put_status(200)
    |> json(%{"predictions" => predictions})
  end

  # バイナリのまま受ける場合
  def index(conn, %{}) do
    {:ok, binary, _} = Plug.Conn.read_body(conn)

    predictions = predict(binary)

    conn
    |> put_status(200)
    |> json(%{"predictions" => predictions})
  end

  # 推論処理
  defp predict(binary) do
    # 画像バイナリを Nx テンソル化
    tensor =
      binary
      |> StbImage.read_binary!()
      |> StbImage.to_nx()

    # 準備しておいた ResNet を取得
    resnet = Serving.get(:resnet)

    # ResNet にテンソルを渡して推論実行
    resnet
    |> Nx.Serving.run(tensor)
    |> then(& &1.predictions)
  end
end

BASE64 文字列とバイナリそのままの両方で受けられる様にしています

router.ex

lib/api_web/router.ex の内容を変更します

...
- scope "/api", ApiWeb do
+ scope "/", ApiWeb do
    pipe_through :api
+   get "/ping", PingController, :ping
+   post "/invocations", PredictionController, :index
  end
end

以下の様にルーティングしています

• /ping の GET で PingController の ping 関数を呼ぶ
• /invocations の POST で PredictionController の index 関数を呼ぶ

ポート番号の 8080 と、このルーティングの仕様は SageMaker で定められています

Phoenix の起動

ターミナルで api ディレクトリーに移動します

cd api

依存モジュールを取得します

mix deps.get

Phoenix を起動します

mix phx.server

しばらくすると Phoenix が起動します

...
[info] TfrtCpuClient created.
[info] Running ApiWeb.Endpoint with cowboy 2.9.0 at 0.0.0.0:8080 (http)
[info] Access ApiWeb.Endpoint at http://localhost:8080

初回時には ResNet のモデルファイルダウンロードが動くので少し時間がかかります

動作確認

別のターミナルを開きます

まず ping を確認します

$ curl http://localhost:8080/ping
{}%

空の JSON が返ってきます

次に推論を確認します

@sample.jpg の部分は手元の JPEG ファイルのパスに変えてください

$ curl -XPOST http://localhost:8080/invocations \
  --data-binary @sample.jpg \
  --header "Content-Type:image/jpeg"
{"predictions":[{"label":"notebook, notebook computer","score":0.6154251098632812},{"label":"laptop, laptop computer","score":0.11529479920864105},{"label":"bow tie, bow-tie, bowtie","score":0.03906528279185295},{"label":"projector","score":0.0267447829246521},{"label":"dining table, board","score":0.024912187829613686}]}%  

推論結果の JSON が返ってきます

jq をインストールしている場合は jq にパイプすると見やすいです

$ curl -XPOST http://localhost:8080/invocations \
  --data-binary @sample.jpg \
  --header "Content-Type:image/jpeg" | jq
{
  "predictions": [
    {
      "label": "notebook, notebook computer",
      "score": 0.6154251098632812
    },
    {
      "label": "laptop, laptop computer",
      "score": 0.11529479920864105
    },
    {
      "label": "bow tie, bow-tie, bowtie",
      "score": 0.03906528279185295
    },
    {
      "label": "projector",
      "score": 0.0267447829246521
    },
    {
      "label": "dining table, board",
      "score": 0.024912187829613686
    }
  ]
}

まとめ

Bumblebee で作成できる Nx.Serving を組み込むだけで、簡単に Phoenix で AI REST API が実装できました

次回はこれをコンテナ化します