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

はじめに

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

実装したコード

実行環境

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

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

コンテナ定義

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

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

#!/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

---

version: '3.2'

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

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

$ 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 してみましたが、、、

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

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

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

コンテナに入ってみて

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 を変更する
Rancher Desktop を再起動する
コンテナを起動する docker-compose up

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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 を忘れずに入れておきます

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 の別名をつけておきます

alias Evision, as: OpenCV

モデルの読込

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

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

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

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

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

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 ができます

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

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

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

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 にリサイズしています

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

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

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

  OpenCV.cvtColor!(OpenCV.cv_COLOR_BGR2RGB())

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

  OpenCV.Nx.to_nx()

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

推論の実行

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

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

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

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

推論結果のラベル表示

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

[452, 434, 808, 515, 431]

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

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

実行結果はこうなります

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

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

推論の関数化

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

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 から適当な画像を見つけてテストしてみましょう

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

まとめ

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

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

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

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

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