複数ノードに接続している場合の Nx.Serving.batched_run の挙動を確認する

## はじめに

前回、1台の端末で機械学習モデルをロードし、もう1台から呼び出してみました

https://qiita.com/RyoWakabayashi/items/9def33deff5ec0f06412

ただ、呼ばれる側が1台だとAIの推論処理自体は分散されていないため、あまり分散らしさを感じませんでした

というわけで、今回は3つのノードで機械学習モデルをロードし、1つのノードから呼び出してみます

今回も Livebook を使っていきます

https://livebook.dev/

:::note warn
そんなにたくさん端末がないので、今回は同一端末内の別ノードです

もちろん、別端末でも同じように動作します
:::

実装したノードブックはこちら

https://github.com/RyoWakabayashi/elixir-learning/tree/main/livebooks/bumblebee/multi_serving

## サーバー側

前回と同じように Vision Transformer のモデルを [Hagging Face](https://huggingface.co/) から読み込みます

以下の処理を3つのノードで実行します

### セットアップ

セットアップセルで依存モジュールをロードし、高速化のためにバックエンドを EXLA に指定します

```elixir
Mix.install(
  [
    {:kino_bumblebee, "~> 0.3.0"},
    {:exla, "~> 0.5.1"}
  ],
  config: [nx: [default_backend: EXLA.Backend]]
)
```

### モデルのロード、子プロセス開始

Vision Transformer を指定して Bumblebee でロードします

`Bumblebee.Vision.image_classification` で画像分類を提供するとき、 `top_k: 1` を指定します

こうすると、画像分類結果が上位1つだけになります

他の2つのノードではここを `top_k: 2` と `top_k: 3` にします

これによって、どのノードで画像分類処理が実行されたのか判別できます

```elixir
{:ok, model_info} = Bumblebee.load_model({:hf, "google/vit-base-patch16-224"})
{:ok, featurizer} = Bumblebee.load_featurizer({:hf, "google/vit-base-patch16-224"})

serving =
  Bumblebee.Vision.image_classification(model_info, featurizer,
    compile: [batch_size: 1],
    defn_options: [compiler: EXLA],
    top_k: 1
  )

Kino.start_child({Nx.Serving, name: ViT, serving: serving})
```

![スクリーンショット 2023-06-11 9.12.39.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/cbf8ff83-8a2d-e154-06ac-3d050c5332cf.png)

### 接続情報取得

クライアント側から接続するためのノード名、クッキーを表示します

```elixir
{node(), Node.get_cookie()}
```

![スクリーンショット 2023-06-11 9.18.26.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/5e135b69-cce3-813e-e3e7-1827507d9152.png)

## クライアント側

### セットアップ

サーバー側と同じモジュールをインストールします

```elixir
Mix.install(
  [
    {:kino_bumblebee, "~> 0.3.0"},
    {:exla, "~> 0.5.1"}
  ],
  config: [nx: [default_backend: EXLA.Backend]]
)
```

### 接続情報の取得

前回は Livebook に直書きしましたが、今回はテキスト入力にします

```elixir
server_node_inputs =
  ["A", "B", "C"]
  |> Enum.into(%{}, fn node_id ->
    {
      node_id,
      %{
        node: Kino.Input.text("SERVER_#{node_id}_NODE_NAME"),
        cookie: Kino.Input.text("SERVER_#{node_id}_COOKIE")
      }
    }
  end)
```

```elixir
server_node_inputs
|> Enum.map(fn {_, inputs} ->
  [inputs.node, inputs.cookie]
end)
|> List.flatten()
|> Kino.Layout.grid(columns: 2)
```

表示されるテキスト入力にそれぞれ3つのノードのノード名、クッキーを入力します

![スクリーンショット 2023-06-11 9.16.40.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/96cf69f5-5523-5d83-ec2c-590c14b21cbe.png)

### サーバー側への接続

テキスト入力を元に、各サーバーに接続します

```elixir
server_node_inputs
|> Enum.map(fn {_, inputs} ->
  node_name =
    inputs.node
    |> Kino.Input.read()
    |> String.to_atom()

  cookie =
    inputs.cookie
    |> Kino.Input.read()
    |> String.to_atom()

  Node.set_cookie(node_name, cookie)

  Node.connect(node_name)
end)
```

実行結果がすべて `true` になれば OK です

```elixir
[true, true, true]
```

### 画像分類の呼び出し

`Nx.Serving.batched_run` で各サーバーの子プロセス `ViT` を呼び出します

```elixir
image_input = Kino.Input.image("Image", size: {224, 224})
form = Kino.Control.form([image: image_input], submit: "Run")
frame = Kino.Frame.new()

Kino.listen(form, fn %{data: %{image: image}} ->
  if image do
    Kino.Frame.render(frame, Kino.Text.new("Running..."))

    image = image.data |> Nx.from_binary(:u8) |> Nx.reshape({image.height, image.width, 3})

    output = Nx.Serving.batched_run(ViT, image)

    output.predictions
    |> Enum.map(&{&1.label, &1.score})
    |> Kino.Bumblebee.ScoredList.new()
    |> then(&Kino.Frame.render(frame, &1))
  end
end)

Kino.Layout.grid([form, frame], boxed: true, gap: 16)
```

表示される画像入力に適当な画像をアップロードし、 `Run` ボタンをクリックします

すると、クリックする度に上位1つ、上位2つ、上位3つの画像分類結果がランダムに返ってきます

![ditributed.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/8e8f9d3d-cc52-882c-b5c1-0a9429bf2b2f.gif)

確かに3つのサーバー側ノードに分散できています

## まとめ

`Nx.batched_run` によって機械学習処理を分散できることを確認しました

次はもっと色々な処理を分散させてみましょう