0.Prologue

「下手な鉄砲も数撃ちゃ…」とはよく言ったものだ。それはまさかの出来事であった。

2022年11月某日、拙作 OnnxInterpの発展的開発終了を記念して、初代TflInterpの開発の発端となった Y系物体検出モデルを一堂に集めてネタのプロジェクト「怒涛のよろず祭り」を作成していた。そんな最中、これらのモデルの中で Axonで読めるモノがあるのかなぁと悪魔の囁きが聞こえ、無駄骨とは思いつつも手を止めて AxonOnnx.importで片っ端から試し始めた。そう、ご存じの通り AxonOnnx.importは堅物で、ここがダメ、あそこがダメと言ってなかなかネットに流布しているonnxモデルを受け付けてくれないのだ。

そしてその瞬間がやって来た…
「ん?? なんか今エラーしなかったような? いやそんな筈はない。見間違いだ」
と、3～4回同じことを繰り返し…
「ええぇ!!! この YOLOv4.onnxのモデル、Axonにコンバートできるやん!!!!!!!!!!」
と我が目を疑ったのである。

という訳で、Axonで読める YOLOv4のモデルが見つかったよぉぉぉ

1.何はともあれ動かしてみる

堅物の AxonOnnx.importの関門を通るのは分かったが、果たしてちゃんと動作するのか? 小生は疑り深い性格である

幸い、「祭り」プロジェクトで動作確認が出来ている YOLOv4.exモジュールとデモ・ドライバYOLOs.livemd[*1]がある。これを AxonInterp向けにチョコチョコっと修正して Livebookでテストしてみよう[*2]。notebookの名前は YOLOv4.livemdでいいかな。

[*1]mix用もあるのだが、やっぱりビジュアルな方が分かり易いよね
[*2]AxonInterpを含む *Interpファミリーは、お互いに実装コードを流用し合えることを目指しています。

YOLOv4.ex

defmodule YOLOv4 do
  @moduledoc """
  Original work:
    Pytorch-YOLOv4 - https://github.com/Tianxiaomo/pytorch-YOLOv4
  """

  use AxonInterp,
    model: "./model/yolov4_1_3_608_608_static.onnx",
    url: "https://github.com/shoz-f/axon_interp/releases/download/0.0.1/yolov4_1_3_608_608_static.onnx",
    inputs: [f32: {1, 3, 608, 608}],
    outputs: [f32: {1, 22743, 1, 4}, f32: {1, 22743, 80}]

  @prepro CImg.builder()
          |> CImg.resize({608,608})
          |> CImg.to_binary([{:range, {0.0, 1.0}}, :nchw])

  def apply(img) do
    # preprocess
    input0 = CImg.run(@prepro, img)
 
    # prediction
    outputs = session()
      |> AxonInterp.set_input_tensor(0, input0, [:binary])
      |> AxonInterp.invoke()

    # postprocess
    boxes = AxonInterp.get_output_tensor(outputs, 0, [:binary])
    scores = AxonInterp.get_output_tensor(outputs, 1, [:binary])

    PostDNN.non_max_suppression_multi_class(
      {22743, 80}, boxes, scores,
      boxrepr: :corner,
      label: "./model/coco.label"
    )
  end
end

LIveYOLOv4.ex

defmodule LiveYOLOv4 do
  @palette CImg.Util.rand_palette("./model/coco.label")

  def run(path) do
    img = CImg.load(path)

    with {:ok, res} <- YOLOv4.apply(img) do
      IO.inspect(res)

      Enum.reduce(res, CImg.builder(img), &draw_item(&1, &2))
      |> CImg.display_kino(:jpeg)
    end
  end

  defp draw_item({name, boxes}, canvas) do
    color = @palette[name]

    Enum.reduce(boxes, canvas, fn [_score, x1, y1, x2, y2, _index], canvas ->
      [x1, y1, x2, y2] = PostDNN.clamp([x1, y1, x2, y2], {0.0, 1.0})

      CImg.fill_rect(canvas, x1, y1, x2, y2, color, 0.35)
    end)
  end
end

そうそう、Axonは EXLAを援用しないと死ぬほど遅いので、setupセルの Mix.installに configオプションを付けておこう

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

Mix.install(
  [
    {:axon_interp, "~> 0.1.0"},
    {:cimg, "~> 0.1.14"},
    {:postdnn, "~> 0.1.5"},
    {:kino, "~> 0.7.0"}
  ],
  config: [
    nx: [default_defn_options: [compiler: EXLA]]
  ]
  # system_env: [{"NNINTERP", "Axon"}]
)

話は前後するが、YOLOv4.livemdを実行するディレクトリの構造は下記の様にした。今回発見した YOLOv4の onnxモデルyolov4_1_3_608_608_static.onnxとcoco.labelはサブディレクトリmodelに置く[*3]。それと、テスト用の画像はいつものdog.jpgでいいよね。

[*3]尤も、上のYOLOv4.exでは、所定のモデル・ファイルが見つからなければ指定のURLからダウンロードせよとAxonInterpに指示していたりする

├── YOLOv4.livemd
├── dog.jpg
└── model
    ├── coco.label
    └── yolov4_1_3_608_608_static.onnx

さぁ準備が整った。テストだ、テスト。わくわくする
モジュール YOLOv4の実体は GenServerなので一度だけ start_link/1を実行。

YOLOv4.start_link([])

そいでもってラン♬

LiveYOLOv4.run("dog.jpg")

おおおおぉ!
ちゃんと動くではないか。
ここにきて、やっと納得した小生であった

PS. 何故か Axonに変換した YOLOv4モデルはシリアライズ、逆コンバートが出来ないようで。一方通行?
どうなってんだこれぇ?

2.オリジナル・ワーク

さて、無事に Axonで YOLOv4が動いた訳ですが、いったいこの pre-traindモデルを何処から拾ってきたモノのか気になりますよね。本家本元の Darknetプロジェクト……ではなく、それの Pytorchへの移植を行っている下記のプロジェクトから拾ってきたのでした。先人達のお仕事に感謝・感謝です

でもって、onnxモデルの生成は、プロジェクトに記述されていた手順を元に google Colabotryで下記の通り行いました。出来上がった onnxモデルはオリジナルの Darknet版をコンバートしたモノの様です。利便性を考えて、この onnxモデルを小生の GitHub ならびに Google Driveで公開中です。

GitHub:
-- https://github.com/shoz-f/axon_interp/releases/download/0.0.1/yolov4_1_3_608_608_static.onnx
Google Drive:
-- https://drive.google.com/file/d/1oY9Pv4Q_MfPolG4sRhydf1GGFnv7556c/view?usp=share_link

setup

!git clone https://github.com/Tianxiaomo/pytorch-YOLOv4
%cd pytorch-YOLOv4
!pip install onnx onnxruntime

download darknet weights

!wget https://github.com/AlexeyAB/darknet/releases/download/darknet_yolo_v3_optimal/yolov4.weights

convert to onnx

!python demo_darknet2onnx.py cfg/yolov4.cfg data/coco.names yolov4.weights data/dog.jpg 1

モデルの入力は一つ、出力は二つで、それぞれ次の仕様となっています。特筆すべき点は、このモデルの出力はそのまま後処理の Multi-class NMSに渡せるという点です(YOLOv4.ex参照)。これは、DNN推論プログラムの実装者にとってとても有り難いことです

入力1:
RGB画像を{608,608}にリサイズし、各画素を{0.0,1.0}の範囲のfloat32に正規化、NCHW変換した tensor

出力1:
float32[1][22743][1][4]の tensorで、22743個のバウンディング・ボックスの対角点座標(x1,y1)-(x2,y2)を{608,608}に対する比率の表したもの。

出力2:
float32[1][22743][80]の tensorで、22743個の各バウンディング・ボックスが cocoの80カテゴリの各々オブジェクトで有り得る確率を表したもの。

余談ですが、YOLOv4が発表された頃には、Object Detectionモデルはシステマティックなアーキテクチャが主流となっており、それまでに蓄積された知見や技術を組み合わせ、より高精度かつ高速性を追い求める時代となっていたようです。YOLOv4も多分に漏れず、下記の主流のアーキテクチャで、特徴マップを作るBackboneに"CSPDarknet53"を、オブジェクトのクラスとバウンディング・ボックスを推定するHeadに"YOLOv3"を、そしてBackbornの各層(解像度?)の情報を混合してHeadに渡すNeckに"SPP"/"PAN"を配しているとのことです。それ以外にも、ユニット/オペレーションレベルでの工夫もされているようです(詳しくは理解していない)。