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

はじめに

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

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

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

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

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

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

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

サーバー側

前回と同じように Vision Transformer のモデルを Hagging Face から読み込みます

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

セットアップ

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

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 にします

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

{: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})

接続情報取得

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

{node(), Node.get_cookie()}

クライアント側

セットアップ

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

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

接続情報の取得

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

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)

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

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

サーバー側への接続

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

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 です

[true, true, true]

画像分類の呼び出し

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

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つの画像分類結果がランダムに返ってきます

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

まとめ

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

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