2024/06/21 現在 Axon を使う場合の注意（GitHubから最新を取得すること）

## はじめに

久しぶりに Livebook で Axon を使ってみたら上手く動かなかったので共有します

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

https://github.com/RyoWakabayashi/elixir-learning/blob/main/livebooks/axon/or_train.livemd

## 実行したコード

動かしたのは以下の記事で実行していたコードです

OR 演算を Axon で学習しています

https://qiita.com/RyoWakabayashi/items/62e45b805f5de2d96017

## モジュールのインストール

2024/06/21 現在、 Axon の最新版は 0.6.1 になっています

https://hex.pm/packages/axon

以下のように各モジュール最新版を使うように指定します

```elixir
Mix.install([
  {:nx, "~> 0.7"},
  {:axon, "~> 0.6"},
  {:exla, "~> 0.7"},
  {:kino, "~> 0.13"},
  {:kino_vega_lite, "~> 0.1"}
])
```

インストールされるものは以下の通りです

```
  axon 0.6.1
  complex 0.5.0
  elixir_make 0.8.4
  exla 0.7.2
  fss 0.1.1
  kino 0.13.0
  kino_vega_lite 0.1.11
  nimble_pool 1.1.0
  nx 0.7.2
  polaris 0.1.0
  table 0.1.2
  telemetry 1.2.1
  vega_lite 0.1.9
  xla 0.6.0
```

Axon のマクロを使うため、 `require Axon` しておきます

```elixir
require Axon
```

## 学習データの生成

元の記事と同じコードです

```elixir
generate_train_data = fn ->
  inputs =
    1..2
    |> Enum.into(%{}, fn index ->
      {
        "input#{index}",
        1..32
        |> Enum.map(fn _ -> Enum.random(0..1) end)
        |> Nx.tensor()
        |> Nx.new_axis(1)
      }
    end)

  labels = Nx.logical_or(inputs["input1"], inputs["input2"])

  {inputs, labels}
end
```

```elixir
generate_train_data.()
```

```elixir
train_data =
  generate_train_data
  |> Stream.repeatedly()
  |> Enum.take(1000)
```

これで、ランダムな `0` `1` の入力二つに対して、 OR 演算した結果の出力という組み合わせが 1000 個作成できました

## モデル定義

ここも元のコードと同じです

```elixir
input1 = Axon.input("input1", shape: {nil, 1})
input2 = Axon.input("input2", shape: {nil, 1})
```

```elixir
model =
  Axon.concatenate(input1, input2)
  |> Axon.dense(8, activation: :relu)
  |> Axon.dense(1, activation: :sigmoid)
```

## トレーニング

損失と正解率のグラフ出力エリアを用意します

```elixir
loss_plot =
  VegaLite.new(width: 300)
  |> VegaLite.mark(:line)
  |> VegaLite.encode_field(:x, "step", type: :quantitative)
  |> VegaLite.encode_field(:y, "loss", type: :quantitative)
  |> Kino.VegaLite.new()

acc_plot =
  VegaLite.new(width: 300)
  |> VegaLite.mark(:line)
  |> VegaLite.encode_field(:x, "step", type: :quantitative)
  |> VegaLite.encode_field(:y, "accuracy", type: :quantitative)
  |> Kino.VegaLite.new()

Kino.Layout.grid([loss_plot, acc_plot], columns: 2)
```

トレーニングを実行します

```elixir
trained_state =
  model
  |> Axon.Loop.trainer(:binary_cross_entropy, :sgd)
  |> Axon.Loop.metric(:accuracy, "accuracy")
  |> Axon.Loop.kino_vega_lite_plot(loss_plot, "loss", event: :epoch_completed)
  |> Axon.Loop.kino_vega_lite_plot(acc_plot, "accuracy", event: :epoch_completed)
  |> Axon.Loop.run(train_data, %{}, epochs: 5, iterations: 1000, compiler: EXLA)
```

これを実行すると、問題が発生します

## トレーニングで損失が増加する

以前の実行結果（正しいもの）は以下のように損失が減少し、正解率が上昇しています

- 損失

  ![loss_graph.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/20d02ea4-b8d5-6264-2c96-4e5f2c747000.png)

- 正解率

  ![acc_graph.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/8336b7c9-772e-55a8-5acb-dc4078e92e45.png)

ところが、 Axon 0.6.1 を使うと以下のように損失が上昇し、正解率が減少しています

- 損失

  ![LOSS](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/ac526532-2916-19c8-5c75-d4aeeff7005c.png)

- 正解率

  ![ACC](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/d3201aef-84cb-4801-1c2e-cefd381eb2f3.png)

これでは全くトレーニングできていません

## 対処

この現象は Axon 0.6.1 のバグです

GitHub の Issue にも登録されていて 2024/05/14 に解消済ですが、反映されたバージョンがまだ Hex にリリースされていません

https://github.com/elixir-nx/axon/issues/567

そのため、 Axon は GitHub から最新版を取得する必要があります

先頭のセットアップセルを以下のように書き換えてください

```elixir
Mix.install([
  {:nx, "~> 0.7"},
  {:axon, "~> 0.6", git: "https://github.com/elixir-nx/axon/"},
  {:exla, "~> 0.7"},
  {:kino, "~> 0.13"},
  {:kino_vega_lite, "~> 0.1"}
])
```

差分

```diff_elixir
-  {:axon, "~> 0.6"},
+  {:axon, "~> 0.6", git: "https://github.com/elixir-nx/axon/"},
```

また、トレーニングのコードも以下のように変更してください

初期状態として `Axon.ModelState.empty()` を渡すようになっています（ `%{}` でも警告されるだけでエラーにはならない）

```elixir
trained_state =
  model
  |> Axon.Loop.trainer(:binary_cross_entropy, :sgd)
  |> Axon.Loop.metric(:accuracy, "accuracy")
  |> Axon.Loop.kino_vega_lite_plot(loss_plot, "loss", event: :epoch_completed)
  |> Axon.Loop.kino_vega_lite_plot(acc_plot, "accuracy", event: :epoch_completed)
  |> Axon.Loop.run(train_data, Axon.ModelState.empty(), epochs: 5, iterations: 1000, compiler: EXLA)
```

差分

```diff_elixir
-  |> Axon.Loop.run(train_data, %{}, epochs: 5, iterations: 1000, compiler: EXLA)
+  |> Axon.Loop.run(train_data, Axon.ModelState.empty(), epochs: 5, iterations: 1000, compiler: EXLA)
```

この状態で実行すると、以前と同じように正しくトレーニングできます

## Explorer と併用する際の注意

以下のように Axon の GitHub 最新版と Explorer の Hex 最新版を併用しようとすると、依存関係のエラーが発生します

```elixir
Mix.install([
  {:exla, "~> 0.7"},
  {:axon, "~> 0.6", git: "https://github.com/elixir-nx/axon/"},
  {:kino, "~> 0.13"},
  {:kino_vega_lite, "~> 0.1"},
  {:explorer, "~> 0.8"}
])
```

エラー

```
Unchecked dependencies for environment dev:
* table_rex (Hex package)
  the dependency table_rex 4.0.0

  > In deps/explorer/mix.exs:
    {:table_rex, "~> 3.1.1 or ~> 4.0.0", [env: :prod, hex: "table_rex", repo: "hexpm", optional: false]}

  does not match the requirement specified

  > In deps/axon/mix.exs:
    {:table_rex, "~> 3.1.1", [env: :prod, hex: "table_rex", optional: true, repo: "hexpm"]}

  Ensure they match or specify one of the above in your deps and set "override: true"
```

エラーメッセージにある通り、 `"override: true"` を指定して table_rex をインストールすればエラーが解消されます

```elixir
Mix.install([
  {:exla, "~> 0.7"},
  {:axon, "~> 0.6", git: "https://github.com/elixir-nx/axon/"},
  {:kino, "~> 0.13"},
  {:kino_vega_lite, "~> 0.1"},
  {:explorer, "~> 0.8"},
  {:table_rex, "~> 4.0", override: true}
])
```

## まとめ

モジュールのバージョンが上がったときは注意が必要です

GitHub の Issues などを見に行き、バグが報告されていないか確認しましょう