MagikaのONNXモデルを動かしてみる

Google製のファイルフォーマット判別ツールMagikaの1.0が公開されたというニュースが今月ありました。
https://www.publickey1.jp/blog/25/googleaimagika_10rust200.html

Magikaはファイルフォーマットの推測にONNXモデルを使っているので、機械学習もまともにやったことが無い人間だけど、折角だからこれで入門してみようと思います。

モデルの入手

MagikaのモデルはGitHubで公開されているリポジトリーのassets/modelsディレクトリーにあります。

最新のstandard_v3_3ディレクトリーからmodel.onnxをダウンロードします。

入力と出力の確認

モデルの入力と出力の型やシェイプを確認します。

learn-magika.rb

require "onnxruntime"

session = OnnxRuntime::InferenceSession.new("model.onnx")
pp session

% ruby ./learn-magika.rb
#<OnnxRuntime::InferenceSession:0x0000000105099330
 @allocator=#<FFI::Pointer address=0x000000012844f990>,
 @inputs=[{name: "bytes", type: "tensor(int32)", shape: ["unk__214", 2048]}],
 @outputs=[{name: "target_label", type: "tensor(float)", shape: ["unk__215", 214]}],
 @session=#<FFI::AutoPointer address=0x0000000711598000>>

入力（@inputs）にも出力（@outputs）にも不明（"unk__..."）な次元がありますね……。RustやPythonでモデルを触っている箇所を見てみたところ

• 可変長配列
• 各エントリーはファイルに対応
• 配列の各要素は長さ2048の整数（0から255）の配列（ファイル内容の部分バイト列）
• 前半1024バイトはファイルの冒頭のバイト列
• 後半1024バイトはファイルの末尾のバイト列
• 足りない分は256で埋める
• 出力は入力と同じ数だけの配列

ということのようです。

入力の準備

引数に渡されたファイルパスから内容を読み込んで、モデルへの入力を準備します。
ここではファイルは充分に大きい、つまり2048バイト以上だと想定します。

learn-magika.rb

input = Array.new(2048, 256)
File.open ARGV[0], "rb" do |file|
  input[...1024] = file.read(1024).bytes
  file.seek(-1024, IO::SEEK_END)
  input[1024..] = file.read(1024).bytes
end

実際にはファイルが小さい時に256で埋める処理とか、空白を処理したりとかした方がいいようですが、一先ず動かすところまで行きたいので省略します。

推論（判定）

いよいよ推論（判定）です。
出力を確認すると"target_label"という名前で出力されるようなのでこれを取得します。

learn-magika.rb

raw_preds = session.run(["target_label"], {bytes: [input]})
pp raw_preds
pp raw_preds[0][0].length

今回は一つのファイルなので、一要素配列（[input]）を渡しています。複数ファイルの場合はここを複数にします。

実行。

% ruby ./learn-magika.rb
[[[2.462988035745184e-09,
   1.2777570645994274e-07,
   1.9437348441897484e-07,
   
   (snip)
   
   1.4691248217957309e-09,
   1.4292183436737105e-07,
   1.4689343075247052e-09]]]
214

なんかでかい配列が返ってきました。出力の定義にshape: ["unk__215", 214]とあるので、これと合致していますね。

ラベル付け

単に配列が返ってきても困ってしまいますが、各要素は、予め決められたファイルフォーマット一覧に対応して、そのフォーマットである確率を表しているみたいです。そう言えば機械学習ってそういう出力の仕方をするのでしたね、何かで読んだ記憶があります。

実際、合計してみると

learn-magika.rb

pp raw_preds[0][0].sum

% ruby ./learn-magika.rb
1.0000002069796026

おおよそ1になるのでちゃんと確率になっていそう。

で、各要素がどのフォーマットに対応しているのか、という情報ですが……
モデルをダウンロードしたディレクトリーにありました。
ここのconfig.min.jsonファイルのtarget_labels_spaceがフォーマットの一覧です。

対応させてみましょう。

learn-magika.rb

require "json"
config = JSON.load_file(CONFIG_PATH, symbolize_names: true)
labels = config[:target_labels_space]
probs = raw_preds[0][0]
result = labels.zip(probs).to_h
pp result

% ruby ./learn-magika.rb model.onnx
{"3gp" => 2.462988035745184e-09,
 "ace" => 1.2777570645994274e-07,
 "ai" => 1.9437348441897484e-07,
 
 (snip)
 
 "onnx" => 0.9999382495880127,
 
 (snip)
 
 "zig" => 1.4691248217957309e-09,
 "zip" => 1.4292183436737105e-07,
 "zlibstream" => 1.4689343075247052e-09}

ちゃんとONNXモデルであることが判定できているようですね。99%以上の確率！

Magikaは、ファイルによって判定を信じる閾値が異なるようで、デフォルトは0.5ですが、TSVは0.9以上じゃないと信用しないとかあるようです。今回はモデルを動かすところまでにしますが、これもconfig.min.jsonにあるので、アプリケーションに組み込む際は加味して判定する必要があります。

Numo::NArrayの導入

さて、ここで

1. ファイルから読み込んだバイト列をRubyの文字列にして
2. それを数値配列にして、
3. それをまたCの数値配列にして、
4. それをONNXモデルで判定する、

というのは非効率です（一ファイルなら気にならないような気もしますが、そういうことにしておいてください）。

RubyのONNX Runtimeは「Ruby界のnumpy」ことNumo::NArray入力にも対応しているので、これを試してみたいと思います。初Numoです。

learn-magika-numo.rb

require "numo/narray"
require "onnxruntime"
require "json"

CONFIG_PATH = File.expand_path("~/src/github.com/google/magika/python/src/magika/models/standard_v3_3/config.min.json")

session = OnnxRuntime::InferenceSession.new("model.onnx")
config = JSON.load_file(CONFIG_PATH, symbolize_names: true)

beginning_feature, ending_feature = nil, nil
File.open ARGV[0], "r" do |file|
  beginning_feature = Numo::Int8.from_binary(file.read(1024))
  file.seek(-1024, IO::SEEK_END)
  ending_feature = Numo::Int8.from_binary(file.read(1024))
end
input = beginning_feature
          .append(ending_feature)
          .cast_to(Numo::Int32)
          .reshape(1, 2048)

raw_preds = session.run(["target_label"], {bytes: OnnxRuntime::OrtValue.from_numo(input)})

labels = config[:target_labels_space]
probs = raw_preds[0][0]
result = labels.zip(probs).to_h
pp result

（出力は同じなので省略）

正直、コマンドラインからの実行で一ファイルだけ判定するというのだと、ちょっと速くなったかな程度ですが、とにかくNumoを使ってもファイルフォーマット判定をすることができました。

これを元にgemも作れるかなあ。
複数ファイル読み込みの高速化とかされているRust実装のバインディングを作った方が実用的な気もしますが、ONNXランタイムを直接実行してみるというのもツールボックスに入れておくと楽しい気がしました。

おまけ：GPUを使う

GPUを使うこともできます。

ONNX Runtimeのリリースページから自分の環境向けのファイルをダウンロード・展開した後、

learn-magika.rb

OnnxRuntime.ffi_lib = "./onnxruntime-osx-arm64-1.23.2/lib/libonnxruntime.dylib"
session = OnnxRuntime::InferenceSession.new("model.onnx", providers: ["CoreMLExecutionProvider"])

でCoreMLを使うようになる……筈がエラーで終了してしまいますね。

% ruby ./learn-magika-numo.rb model.onnx
libc++abi: terminating due to uncaught exception of type onnxruntime::OnnxRuntimeException: /Users/runner/work/1/s/include/onnxruntime/core/common/logging/logging.h:371 static const Logger &onnxruntime::logging::LoggingManager::DefaultLogger() Attempt to use DefaultLogger but none has been registered.

[1]    46429 abort      ruby ./learn-magika-numo.rb model.onnx

分かる方、教えてください……。