AxonOnnx を使って VGG16 を Livebook で動かしてみた

## はじめに

この記事は @the_haigo さんの記事を参考に、
Livebook で Axon を動かしてみた際の感想、備忘録です

https://qiita.com/the_haigo/items/8f5157a185e08f6d6bce

実装したコード

https://github.com/RyoWakabayashi/axon-notebooks

## 実行環境

Elixir で AI をやってみるにあたって、色々試行錯誤するだろうということで、
Docker + Livebook で動かしてみる事にしました

- OS: macOS Monterey 12.4
- Docker: 20.10.16
- Rancher Desktop: 1.4.1

## コンテナ定義

参考にした記事では Pytorch の VGG16 を ONNX に変換する際、
Python も使っていたので Jupyter も入れて並列起動させました

Dockerfile

```Dockerfile
FROM livebook/livebook

# 何かと便利なものたち
# Evision も使いたいので OpenCV 関連
RUN apt upgrade -y \
  && apt update \
  && apt install --no-install-recommends -y \
  gnupg2 \
  apt-transport-https \
  libopencv-dev \
  build-essential \
  erlang-dev \
  software-properties-common \
  sudo \
  && apt clean \
  && rm -rf /var/lib/apt/lists/*

# Python3.9 を python 、 pip で使う
RUN apt update \
  && apt install --no-install-recommends -y \
  python3.9 \
  python3-pip \
  && apt clean \
  && rm -rf /var/lib/apt/lists/* \
  && update-alternatives --install /usr/bin/python python /usr/bin/python3.9 10 \
  && update-alternatives --install /usr/bin/pip pip /usr/bin/pip3 10

# Python の依存ライブラリたち
RUN pip install --upgrade pip \
  && pip install jupyterlab torch torchvision

ENV HOME=/home/livebook
# Evision はコンパイル済のものを使う（ビルドにすごく時間がかかるので）
ENV EVISION_PREFER_PRECOMPILED=true

WORKDIR /home/livebook

# Elixir の依存ライブラリをインストールするための準備
RUN mix local.hex --force \
  && mix local.rebar --force

# Jupyter と Livebook を並列で起動するシェルスクリプト
COPY ./run_servers.sh /root/run_servers.sh

RUN chmod +x /root/run_servers.sh

# ノートブックは出来上がったものをコンテナ内にコピー
# volumes でマウントすると、 Elixir の依存ライブラリのインストールでエラーになった
COPY ./notebooks /home/livebook/notebooks

CMD ["/root/run_servers.sh"]
```

run_servers.sh

```bash
#!/bin/bash

# 末尾に & を付けると並列実行
/app/bin/livebook start &

jupyter lab --allow-root --ip=0.0.0.0 --no-browser --notebook-dir=/home/livebook &

wait -n

exit $?
```

docker-compose.yml

```yml
---

version: '3.2'

services:
  livebook:
    build: .
    container_name: livebook
    ports:
      - '8080:8080'
      - '8888:8888'
    volumes:
      - ./data:/data
```

`docker-compose up` で起動すれば、
Jupyter と Livebook 両方の URL が表示されます

```bash
$ docker-compose up --build         
[+] Building 0.2s (14/14) FINISHED                                                                                                                                    
...
[+] Running 1/1
 ⠿ Container livebook  Recreated                                                                                                                                 0.7s
Attaching to livebook
...
livebook  | [C 2022-06-13 08:31:17.012 ServerApp] 
livebook  |     
livebook  |     To access the server, open this file in a browser:
livebook  |         file:///home/livebook/.local/share/jupyter/runtime/jpserver-8-open.html
livebook  |     Or copy and paste one of these URLs:
livebook  |         http://xxx:8888/lab?token=...
livebook  |      or http://127.0.0.1:8888/lab?token=...
livebook  | [Livebook] Application running at http://localhost:8080/?token=...
```

## Pytorch から ONNX への変換

ここは元記事 https://qiita.com/the_haigo/items/8f5157a185e08f6d6bce#pytorch-onnx-export のままです

出力先だけ `/data/vgg16.onnx` にしています

## ONNX から　dets への変換

ここが最大の難所でした

元記事のとおり、 AxonOnnx をインストールして、 import してみましたが、、、

```elixir
Mix.install([
  {:axon_onnx, github: "elixir-nx/axon_onnx"}
])
```

```elixir
{model, params} = AxonOnnx.import("/data/vgg16.onnx")
```

`AxonOnnx.import` を実行すると、しばらくしてプロセスが kill されてしまい、
処理が実行できませんでした

コンテナに入ってみて

```bash
docker exec -it livebook /bin/bash
```

`top` コマンドでメモリを確認してみると、
`AxonOnnx.import` の実行中、明らかにメモリがどんどん使用され、
100% に達するところでプロセスが落ちていました

VGG16 だとメモリ消費はしょうがないのか、、、？

私は macOS で Rancher Desktop を使って Docker を動かしていたので、
Rancher Desktop の設定を変更してみます

変更の手順

- 現在のコンテナを停止する（`docker-compose up` したターミナルで control + c）
- 現在のコンテナを削除する `docker-compose down`
- Rancher Desktop の Kubernetes Settings から Memory を変更する

  ![スクリーンショット 2022-06-13 17.27.40.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/1da01c29-0a8b-d241-5b9d-f017ef8289b7.png)

- Rancher Desktop を再起動する
- コンテナを起動する `docker-compose up`

コンテナ内で `free` や `top` を実行してメモリの total が変わっていれば OK です

ちょっとずつ増やしていった結果、メモリを 10GB まで増やすと `AxonOnnx.import` が実行できました

、、、本当にこんなにメモリが必要？

## EXLA を使ってみる（ダメでした）

Axon は Nx を使って行列計算しています

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

https://github.com/elixir-nx/nx/tree/main/nx#readme

そして、 Nx が実際に行列計算を行う際、バックエンドを選択できるようになっています

デフォルトではバイナリバックエンドが使われます

これは Elixir のコードで一生懸命計算するバックエンドです

https://github.com/elixir-nx/nx/blob/main/nx/lib/nx/binary_backend.ex

残念ながらこれは非常に遅く、メモリ効率も悪いです

Python で行列計算を行う際、 Array のまま各要素毎に + や - で演算するようなものです

当然、 Python であれば numpy を使わなければいけません

Nx の場合も、より速くメモリ効率の良いバックエンドを使わなければなりません

というわけで、 EXLA を使ってみます

https://github.com/elixir-nx/nx/tree/main/exla

EXLA は Google XLA という高速演算コンパイラを Nx のバックエンドとして提供してくれます

インストール対象に EXLA を追加します

```elixir
Mix.install([
  {:axon_onnx, github: "elixir-nx/axon_onnx"},
  {:exla, "~> 0.2.1"}
])
```

そして、 EXLA をバックエンドとして使用するように設定します

```
EXLA.set_as_nx_default([:tpu, :cuda, :rocm, :host])
```

この状態で実行すると、 メモリ 4GB でも落ちることなく、なおかつ速く変換できました

、、、が、

一見正常終了したように見せて、変換後のモデルを推論に使ってみるとエラーに、、、

もしかしたら私の環境等の問題かもしれませんが、、、

現状では メモリをたっぷり確保してバイナリバックエンドを使うしかなさそうです

## Axon による推論

元記事では stb_image を使っていましたが、 OpenCV　に慣れ親しんでいるので Evision を使いましょう

推論時はちゃんと動くので、 EXLA を忘れずに入れておきます

https://github.com/cocoa-xu/evision

```elixir
Mix.install([
  {:download, "~> 0.0.4"},
  {:evision, "~> 0.1.0-dev", github: "cocoa-xu/evision", branch: "main"},
  {:kino, "~> 0.5.2"},
  {:nx, "~> 0.1", [env: :prod, repo: "hexpm", hex: "nx", optional: true]},
  {:exla, "~> 0.2.1"},
  {:axon, "~> 0.1.0-dev", github: "elixir-nx/axon", branch: "main"}
])
```

OpenCV として使いたいので、 Evision には OpenCV の別名をつけておきます

```elixir
alias Evision, as: OpenCV
```

### モデルの読込

dets に変換したモデルを読み込みます

```elixir
{:ok, params} = :dets.open_file('/data/vgg16.dets')
[{1, {model, params}}] = :dets.lookup(params, 1)
```

### クラス番号とラベルの変換用マップ

ここにある ImageNet のクラス一覧を読み込んで、
クラスの番号とラベルの Map を作っておきます

https://gist.github.com/yrevar/942d3a0ac09ec9e5eb3a

もっとスマートに書けると思いますが、、、

```elixir
File.rm("imagenet1000_clsidx_to_labels.txt")

class_list =
  Download.from(
    "https://gist.githubusercontent.com/yrevar/942d3a0ac09ec9e5eb3a/raw/238f720ff059c1f82f368259d1ca4ffa5dd8f9f5/imagenet1000_clsidx_to_labels.txt"
  )
  |> elem(1)

class_map =
  class_list
  |> File.read!()
  |> String.split("\n")
  |> Enum.reduce(%{}, fn line, acc ->
    [class_number, class_name] = String.split(line, ":")

    class_number =
      class_number
      |> String.replace("{", "")
      |> String.trim()
      |> Integer.parse()
      |> elem(0)

    class_name =
      class_name
      |> String.replace("}", "")
      |> String.trim()
      |> String.replace(", ", "!")
      |> String.replace(",", "")
      |> String.replace("!", ", ")
      |> String.replace("'", "")

    acc |> Map.put(class_number, class_name)
  end)
```

こんな感じの Map ができます

```elixir
%{
  912 => "worm fence, snake fence, snake-rail fence, Virginia fence",
  326 => "lycaenid, lycaenid butterfly",
  33 => "loggerhead, loggerhead turtle, Caretta caretta",
  ...
}
```

### 画像のダウンロード、前処理

いつもの Lenna さんの画像を読み込んで、テンソルにします

![lenna](https://upload.wikimedia.org/wikipedia/en/7/7d/Lenna_%28test_image%29.png)

```elixir
File.rm("Lenna_%28test_image%29.png")

lenna =
  Download.from("https://upload.wikimedia.org/wikipedia/en/7/7d/Lenna_%28test_image%29.png")
  |> elem(1)

mat = OpenCV.imread!(lenna)

tensor =
  OpenCV.resize!(mat, [_width = 224, _height = 224])
  |> OpenCV.cvtColor!(OpenCV.cv_COLOR_BGR2RGB())
  |> OpenCV.Nx.to_nx()
  |> Nx.divide(255)
  |> Nx.subtract(Nx.tensor([0.485, 0.456, 0.406]))
  |> Nx.divide(Nx.tensor([0.229, 0.224, 0.225]))
  |> Nx.transpose()
  |> Nx.new_axis(0)
```

元記事と違うのは以下の箇所です

- Evision の `resize` を使って 224 * 224 にリサイズしています

```elixir
  OpenCV.resize!(mat, [_width = 224, _height = 224])
```

- OpenCV だと色が BGR になっているので `cvtColor` で RGB に変換しています

  この辺りは本当に Python と同じ感覚です

```elixir
  OpenCV.cvtColor!(OpenCV.cv_COLOR_BGR2RGB())
```

- `Nx.to_nx` で Nx で演算できるテンソルに変換しています

```elixir
  OpenCV.Nx.to_nx()
```

その後は元記事と同じようにモデルに入力するための前処理を実行しています

### 推論の実行

`Axon.predict` で推論を実行します

```elixir
preds =
  Axon.predict(model, params, tensor)
  |> Nx.flatten()
  |> Nx.argsort()
  |> Nx.reverse()
  |> Nx.slice([0], [5]) # 先頭5件だけ取り出す
  |> Nx.to_flat_list()
```

ここで EXLA をバックエンドに設定し忘れていると、いつまで経っても結果が返ってきません

必ず EXLA をバックエンドに設定しましょう

### 推論結果のラベル表示

preds には以下のようにクラスの番号が入っています

```elixir
[452, 434, 808, 515, 431]
```

クラスの番号では何のことかわからないので、作っておいた class_map を使ってラベルに変換します

```elixir
preds
|> Enum.map(fn element ->
  Map.get(class_map, element)
end)
```

実行結果はこうなります

```elixir
["bonnet, poke bonnet", "bath towel", "sombrero", "cowboy hat, ten-gallon hat", "bassinet"]
```

ボンネットやバスタオル、バシネットはおかしいですが、
ソンブレロやカウボーイハットは妥当な感じです

### 推論の関数化

画像のダウンロードや推論を一つの関数にまとめます

```elixir
defmodule Detector do
  def detect(image_url, model, params, class_map) do
    basename =
      image_url
      |> URI.parse()
      |> Map.fetch!(:path)
      |> Path.basename()

    File.rm(basename)

    mat =
      image_url
      |> Download.from()
      |> elem(1)
      |> OpenCV.imread!()

    tensor =
      OpenCV.resize!(mat, [_width = 224, _height = 224])
      |> OpenCV.cvtColor!(OpenCV.cv_COLOR_BGR2RGB())
      |> OpenCV.Nx.to_nx()
      |> Nx.divide(255)
      |> Nx.subtract(Nx.tensor([0.485, 0.456, 0.406]))
      |> Nx.divide(Nx.tensor([0.229, 0.224, 0.225]))
      |> Nx.transpose()
      |> Nx.new_axis(0)

    model
    |> Axon.predict(params, tensor)
    |> Nx.flatten()
    |> Nx.argsort()
    |> Nx.reverse()
    |> Nx.slice([0], [5])
    |> Nx.to_flat_list()
    |> Enum.map(fn element ->
      Map.get(class_map, element)
      |> IO.puts()
    end)

    Helper.show_image(mat)
  end
end
```

Open Images Dataset から適当な画像を見つけてテストしてみましょう

![スクリーンショット 2022-06-13 18.42.41.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/1db5af0d-694b-08e0-fc2b-ae62f809e034.png)

![スクリーンショット 2022-06-13 18.43.07.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/16d7e596-dc0e-7855-94c4-05738b9ed29c.png)

![スクリーンショット 2022-06-13 18.48.41.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/e445b061-a74d-9317-6f16-334ed7b5d5f0.png)

![スクリーンショット 2022-06-13 18.49.38.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/e64d5c6f-601f-6c7e-329e-9192c4dc058f.png)

いい感じに検出できているようです

## まとめ

Nx を使う場合、バックエンドが重要であることがよく分かりました

EXLA バックエンドで変換がうまくいけば、メモリ問題は解消するのですが、、、

また、やはり Evision を使うと Python のときと同じ感覚で画像処理できるので、
ほとんど実装方法に迷うことがありません

他のモデルも試してみたいと思います