AxonでResNet18を我流でやってみた【利用者編】

# 0.Prologue
Elixir界隈の進歩からみると今更感が満載なのだが、ちょこっと Axonを使って ResNet18による画像分類(推論)をやってみた。

尤も、天邪鬼な小生のこと故、そのやり方は我流も我流だ。道楽で作っている*Interpシリーズの終点にあたる AxonInterpを、やっつけで実装して試用している。

以下、とにかくエイッやぁ～と Livebook上で ResNet18を動かすやり方を駆け足で紹介する。なお、mixで動かしてみたいと思う方は、下記のGithubを参考にして頂きたい。

> https://github.com/shoz-f/axon_interp/tree/main/demo_resnet18

# A.手ぶらでどうぞ編
コードの中身なんてどうでもいいから、とにかくササっと試してみたい方はこちらをどうぞ。
そうでない方は、次の「B.写経したいぞ編」にお進みください:wink:

お好みの場所に作業用のディレクトリを作ります。(ここではホームに "resnet18"ディレクトリを作成しました)

```sh
$ mkdir resnet18
```

Livebookを起動して、`import notebook`ダイアログを開き、下記URLの notebookをインポートします。

> https://github.com/shoz-f/axon_interp/blob/main/livebook_resnet18/resnet18.livemd

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/14158/806fe5d3-c11d-2cec-d6b9-feef671e8a0b.png)

インポートした notebookを、先ほど作成したディレクトリにお好みの名前で保存します。(ここでは名前を"my_resnet18"としました)

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/14158/2a4f99c4-2ad8-ac17-ab17-a6f05f928b5c.png)

notebookのコード・セルを上から順にダァァァっと実行します。成功すれば下の様に、写真の動物を推論した結果が表示されます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/14158/b48c3b21-a252-3cf9-23b5-b071b90b424a.png)


# B.写経したいぞ編

否、どこの馬の骨が書いたか分からない nootbookを盲目的に実行するなんぞ、プログラマの名折れだと憤慨の方もいらっしゃるでしょう。そんな方は、白紙のnootbookに小生が書いたコードを写経しつつ、解析&実験をしてみては如何でしょう？

以下、コードを提示しつつ簡単な説明を加えたいと思います。

## B-1.ResNet18モデルの調達
学習済みResNet18のモデルは、torchvision.modelsから借用しました。Axonで Pytorchのモデルを利用するには、モデルを一旦 ONNX形式に変換し、さらにそれを Elixirデータ構造に変換しなければなりません。

前者のONNXへの変換は、下記の様にPythonスクリプトで行います。お好みのディレクトリでこのスクリプトを実行すると、サブディレクトリ./dataに ONNX形式のモデル`resnet18.onnx`が保存されます。

> $ python mk_resnet18_onnx.py

```python:mk_resnet18_onnx.py
import os
import torch
import torchvision.models as models

model = models.resnet18(weights='DEFAULT')
model.eval()

dummy_input = torch.randn(1, 3, 224, 224)

os.makedirs("data", exist_ok=True)
torch.onnx.export(model,
   dummy_input,
   "data/resnet18.onnx",
   verbose=False,
   input_names=["input.0"],
   output_names=["output.0"],
   export_params=True
   )
```

後者の Elixirデータ構造への変換は、Axonのコンパニオン・モジュール AxonOnnxを用いて行います。先ほどとと同じディレクトリで、下記の exsスクリプトを実行すると、`resnet18.onnx`が Elixirデータ構造に変換され、さらにそれをシリアライズ化(Erlang external term format)したものが./data下の `resnet18.axon`に保存されます。

注)"Erlang external term format"は Erlang/OTPのバージョンに依存する場合があるようで(詳細は未確認)、異なるバージョンの Erlang/Elixirで作成した.axonを読み込むとエラーが発生しました:sob:

> $ elixir mk_resnet18_axon.exs

```elixir:mk_resnet18_axon.exs
Mix.install([
  {:axon_onnx, "~> 0.3.0"}
])

AxonOnnx.import("data/resnet18.onnx")
|> then(fn {model, params} -> Axon.serialize(model, params) end)
|> (&File.write("data/resnet18.axon", &1)).()

```

## B-2.新規notebookの作成
Livebookを起動する前に、作業用のディレクトリを用意します。「手ぶらでどうぞ編」と同じくホームの下に resnet18ディレクトリを用意しました。

```sh
$ mkdir resnet18
```

Livebookを起動し、`[New notebook]`で空の notebookを作成します。notobookの編集が記録されるように、まず最初に作業ディレクトリに notebookを空のまま保存しておきます。(ここではnotebookの名前を"my_resnet18"としました)

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/14158/e045577f-454e-039e-a988-cf521c36a157.png)


## B-3.Setupセル - 依存モジュールの登録
今回作成した ResNet18アプリは、Nx, EXLA, Kino, CImg, AxonInterpモジュールを利用しています。Setupセルに Mix.installを用いてモジュール依存関係を記述しましょう。これらのモジュールのうち `AxonInterp`は小生の試作モジュールのため、参照先はHexではなくGitHubになります。

また、`defn`で EXLAが有効になるように、Mix.installに configオプションを追加します。

Tips) Setupセルに **"File.cd!(\__DIR__)"** の一行を含めておくと、以降のセルの cwdは notobookを保存したディレクトリとなります。

```elixir
File.cd!(__DIR__)
# for windows JP
# System.shell("chcp 65001")

Mix.install(
  [
    {:nx, "~> 0.4.0"},
    {:exla, "~> 0.4.0"},
    {:cimg, "~> 0.1.13"},
    {:axon_interp, github: "shoz-f/axon_interp"},
    {:kino, "~> 0.7.0"}
  ],
  config: [
    nx: [default_defn_options: [compiler: EXLA]]
  ],
  #system_env: [{"NNINTERP", "Axon"}]
)
```

ResNet18アプリは、ImageNetの1000カテゴリのラベル・セットを必要とします。Setupセルまたはそれに続くCodeセルに下記のコードを記述してダウンロードします。ラベル・セットのファイルは、notebookを保存したディレクトリに "imagenet1000.label"の名前で保存されます。

```elixir
# download "imagenet1000.label" if not exists.
unless File.exists?("./imagenet1000.label") do
  AxonInterp.Util.download(
    "https://github.com/shoz-f/axon_interp/releases/download/0.0.1/imagenet1000.label"
  )
end
```

## B-4.ResNet18推論アプリのモジュール
下のモジュールが ResNet18推論アプリの本体です。

prologueで触れた通り、AxonのAPIを直接呼び出す代わりに自作の薄～～いラッパーモジュールAxonInterpを介して Axonを利用しています。以下、コードを簡単に説明します。

`use AxonInterp`から始まる 5～9行目で DNNモデルの読み込みを行っています。モデル・ファイルの置き場所は`model:`オプションで指定します。また、`url:`オプションに URLを指定すると、手元でモデル・ファイルが見つからなかった場合にそのURLからモデル・ファイルをダウンロードしようと試みます。

11行目の@imagenet1000は、ImageNetの 1000カテゴリのラベルをそのインデックス(行番号)をkeyにして Mapにしたものです。

17行目から始まる `apply/2`関数が ResNet18の推論を実行する関数です。第1引数には CImg型の画像を、第２引数には出力する推論結果の個数を与えます。

18～22行目は前処理で、引数で与えられた画像を {224pix, 224pix}にリサイズし、次に平均値/分散値を用いて各画素値を float32の数値型に正規化、最後に NHWCから NCHWに軸交換しています。

24～30行目は DNNモデルによる推論の実行です。推論結果は float32[1000]の tensorとなります。この tensorのインデックスが先のラベルのインデックスと対応づけられています。また tensorの要素の値は画像がそのカテゴリに属する確からしさを表しています。

32～39行目は後処理で、まず推論結果を softmaxに掛けて正規化された確率分布に変換したのちに、推論の確率が高い候補を取り出しラベルに変換して出力します。

```elixir
 1: defmodule ResNet18 do
 2:   @width 224
 3:   @height 224
 4:
 5:   use AxonInterp,
 6:     model: "./data/resnet18.axon",
 7:     url: "https://github.com/shoz-f/axon_interp/releases/download/0.0.1/resnet18.axon",
 8:     inputs: [f32: {1, 3, @height, @width}],
 9:     outputs: [f32: {1, 1000}]
10:
11:   @imagenet1000 (for item <- File.stream!("./imagenet1000.label") do
12:                   String.trim_trailing(item)
13:                 end)
14:                 |> Enum.with_index(&{&2, &1})
15:                 |> Enum.into(%{})
16:
17:   def apply(img, top \\ 1) do
18:     # preprocess
19:     input0 = CImg.builder(img)
20:       |> CImg.resize({@width, @height})
21:       # mean=[0.485, 0.456, 0.406] and std=[0.229, 0.224, 0.225], NCHW.
22:       |> CImg.to_binary([{:gauss, {{123.7, 58.4}, {116.3, 57.1}, {103.5, 57.4}}}, :nchw])
23:
24:     # prediction
25:     outputs = session()
26:       |> AxonInterp.set_input_tensor(0, input0)
27:       |> AxonInterp.invoke()
28:       |> AxonInterp.get_output_tensor(0)
29:       |> Nx.from_binary(:f32)
30:       |> Nx.reshape({1000})
31:
32:     # postprocess
33:     exp = Nx.exp(outputs)
34:
35:     Nx.divide(exp, Nx.sum(exp))  # softmax
36:     |> Nx.argsort(direction: :desc)
37:     |> Nx.slice([0], [top])
38:     |> Nx.to_flat_list()
39:     |> Enum.map(&@imagenet1000[&1])
40:   end
41: end
```

## B-5.Let's try it - 画像を分類してみよう
`ResNet18`モジュールの実体は GenServerなので、最初に一度だけ start_linkで起動します。

```elixir
ResNet18.start_link([])
```

さあ、ResNet18で画像を分類してみましょう。

```elixir
unless File.exists?("./lion.jpg") do
  AxonInterp.Util.download(
    "https://github.com/shoz-f/axon_interp/releases/download/0.0.1/lion.jpg"
  )
end

img = CImg.load("lion.jpg")
Kino.render(CImg.display_kino(img, :jpeg))

ResNet18.apply(img, 3)
```
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/14158/5deee59f-9ffd-43d1-2aca-17ffa67f564c.png)

第１推論:ライオン - ご名答:grin:。第２推論:雄牛 - ん？？？。第３推論:チャウチャウ - おひっ:laughing:

# ∞.Epilogue

以上、駆け足であったが「AxonでResNet18を我流でやってみた」のやり方を紹介した。

DNN推論アプリであれば基本的に同じ処理の流れとなるので、Axonで利用できるモデルが手に入れば、上の ResNet18モジュールをテンプレートとしてササっ[*1]と試すことが出来る……否、*Interpはそれを目的としたシリーズであったのだ:sunglasses:

[*1]実際のところ、前処理/後処理の実装が大変なので ササっとはいかないのだが…

(おしまい)

---
PS.【利用者編】故に AxonInterpの中身については徹底的に触れないようにした:stuck_out_tongue_winking_eye: