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

## はじめに

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

- [画像分類: ResNet50](https://qiita.com/RyoWakabayashi/items/b9b4799683af5aaa29dd)
- [画像生成: Stable Diffusion](https://qiita.com/RyoWakabayashi/items/8070f6bc372015be6a8d)
- [文章の穴埋め: BERT](https://qiita.com/RyoWakabayashi/items/dcc03a8d19d7320dda8f)
- [文章の判別: BERTweet](https://qiita.com/RyoWakabayashi/items/1f8624d717fb2b1dbdc0)
- [文章の生成: GPT2](https://qiita.com/RyoWakabayashi/items/6a780c2858ec952fa1e1)
- [質疑応答: RoBERTa](https://qiita.com/RyoWakabayashi/items/316d0c24020468d07f18)
- [固有名詞の抽出: bert-base-NER](https://qiita.com/RyoWakabayashi/items/82ac1dddfdddc8a2f79d)

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

公式のサンプルはこちら

https://github.com/elixir-nx/bumblebee/blob/main/examples/phoenix/image_classification.exs

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

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

私の実装はこちら

https://github.com/RyoWakabayashi/elixir-learning/tree/main/sagemaker/serve/api

## 実装環境

- macOS 13.1
- Elixir 1.14.2 OTP 25

## プロジェクトの作成

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

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

```bash
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 から不要な行を削除します

```diff_elixir
...
- # 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 に以下の様に依存モジュールを追加します

```diff_elixir
...
  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 のバックエンドとして動かす（処理を高速化する）ためです

```diff_elixir
...
config :phoenix, :json_library, Jason

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

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

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

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

```diff_elixir
...
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 を準備しておきます

```elixir
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 に以下の内容を追加します

```diff_elixir
  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 応答用のコントローラーです

```elixir
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` で定義した画像識別サービスに画像をテンソル化して渡し、推論結果を返します

```elixir
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 文字列とバイナリそのままの両方で受けられる様にしています

https://qiita.com/RyoWakabayashi/items/93192c5eeed0dea90bfc

### router.ex

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

```diff_elixir
...
- 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 で定められています

https://docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-inference-code.html

## Phoenix の起動

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

```bash
cd api
```

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

```bash
mix deps.get
```

Phoenix を起動します

```bash
mix phx.server
```

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

```bash
...
[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 を確認します

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

空の JSON が返ってきます

次に推論を確認します

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

```bash
$ 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 にパイプすると見やすいです

https://stedolan.github.io/jq/


```bash
$ 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 が実装できました

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