RubyでONNXフォーマットのDNNの推論を可能にするExtensionを開発

この度、RubyでCPUによるDeep Neural Networkの推論を可能にするExtension: menoh-ruby を開発・公開しました。

正確には、ONNXフォーマットで出力されたDNNモデルパラメータを用いて推論を行うC++ベースのライブラリmenoh(初出Instant(解説記事))のRubyラッパーを開発しています。

機能概要

• C++ベースのMenohの性能をほぼそのまま活用できる
  • 任意の層のノード値へのアクセスが可能
  • MKL-DNNベースの高速推論
  • バッチサイズ指定により複数画像に対する同時推論が可能
• 入出力をRubyのデータ型に適切にバインドしている
  • 1バッチあたりの入力はRMagickのImageオブジェクトの配列
  • 推論結果をHash/Array形式で取得可能

Rubyコードの実装例

本家のC++による実装例と同等の処理を、このように書けます。
兎にも角にも実装例をご覧ください。

require 'open-uri'
require 'rmagick'
require 'menoh'

# download dependencies
def download_file(url, output)
  return if File.exist? output
  puts "downloading... #{url}"
  File.open(output, 'wb') do |f_output|
    open(url, 'rb') do |f_input|
      f_output.write f_input.read
    end
  end
end
download_file('https://www.dropbox.com/s/bjfn9kehukpbmcm/VGG16.onnx?dl=1', './data/VGG16.onnx')
download_file('https://raw.githubusercontent.com/HoldenCaulfieldRye/caffe/master/data/ilsvrc12/synset_words.txt', './data/synset_words.txt')
download_file('https://upload.wikimedia.org/wikipedia/commons/5/54/Light_sussex_hen.jpg', './data/Light_sussex_hen.jpg')
download_file('https://upload.wikimedia.org/wikipedia/commons/f/fd/FoS20162016_0625_151036AA_%2827826100631%29.jpg', './data/honda_nsx.jpg')

# load dataset
image_list = [
  './data/Light_sussex_hen.jpg',
  './data/honda_nsx.jpg'
]
input_shape = {
  channel_num: 3,
  width: 224,
  height: 224
}

# load ONNX file
onnx_obj = Menoh::Menoh.new './data/VGG16.onnx'

# onnx variable name
CONV1_1_IN_NAME = '140326425860192'.freeze
FC6_OUT_NAME = '140326200777584'.freeze
SOFTMAX_OUT_NAME = '140326200803680'.freeze

# model options for model
model_opt = {
  backend: 'mkldnn',
  input_layers: [
    {
      name: CONV1_1_IN_NAME,
      dims: [
        image_list.length,
        input_shape[:channel_num],
        input_shape[:height],
        input_shape[:width],
      ]
    }
  ],
  output_layers: [FC6_OUT_NAME, SOFTMAX_OUT_NAME]
}
# make model for inference under 'model_opt'
model = onnx_obj.make_model model_opt

# prepare dataset
image_set = [
  {
    name: CONV1_1_IN_NAME,
    data: image_list.map do |image_filepath|
      image = Magick::Image.read(image_filepath).first
      image = image.resize_to_fill(input_shape[:width], input_shape[:height])
      'BGR'.split('').map do |color|
        image.export_pixels(0, 0, image.columns, image.rows, color).map { |pix| pix / 256 }
      end.flatten
    end.flatten
  }
]

# execute inference
inferenced_results = model.run image_set

# load category definition
categories = File.read('./data/synset_words.txt').split("\n")
TOP_K = 5
layer_result = inferenced_results.find { |x| x[:name] == SOFTMAX_OUT_NAME }
layer_result[:data].zip(image_list).each do |image_result, image_filepath|
  puts "=== Result for #{image_filepath} ==="

  # sort by score
  sorted_result = image_result.zip(categories).sort_by { |x| -x[0] }

  # display result
  sorted_result[0, TOP_K].each do |score, category|
    puts "#{category} : #{score}"
  end
end

実行結果(先頭1件分)

$ ruby ./test.rb 
=== Result for ../data/Light_sussex_hen.jpg ===
n01514859 hen : 0.9054788947105408
n01807496 partridge : 0.03229297325015068
n01514668 cock : 0.031027931720018387
n01797886 ruffed grouse, partridge, Bonasa umbellus : 0.014490845613181591
n01798484 prairie chicken, prairie grouse, prairie fowl : 0.0016611808678135276

開発レポジトリ

本プロジェクトはPreferred Networksの研究成果物として、下記レポジトリにて公開しています。

• menoh-ruby

開発の背景

Python系の推論環境はChainerやTensorFlow、PyTorchなどのフレームワークの充実によって多数存在しましたが、Rubyにおける選択肢はスクラッチで再実装するか、プロセスコールやPyCall等を用いて他の言語に依存する手法が一般的でした。
今回、@okdshinさんが
Menohを開発・公開しました。これをRubyから呼び出せるようにすれば様々な可能性が広がると感じたため、開発に着手しました。

今後の展望

ONNXフォーマットはMicrosoftとFacebookが提唱し、普及を進めているもので、Chainerもサポートを進めている状況です。学習と推論の分離、フレームワーク間の互換性の向上等、様々な可能性を秘めたフォーマットで、今後も目が離せません。

具体的な推論結果例

下記に上記実装例の実行結果を載せます。

Category	Score
n04522168 vase	0.14624682068824768
n02948072 candle, taper, wax light	0.12000349164009094
n03930313 picket fence, paling	0.10246642678976059
n12620546 hip, rose hip, rosehip	0.043764956295490265
n02206856 bee	0.03839555010199547

Category	Score
n03877845 palace	0.9945957064628601
n04486054 triumphal arch	0.0019427633378654718
n03781244 monastery	0.0006744712009094656
n03837869 obelisk	0.00033667951356619596
n02980441 castle	0.0003334969806019217

Category	Score
n03314780 face powder	0.27068012952804565
n04019541 puck, hockey puck	0.1563795506954193
n03937543 pill bottle	0.14620213210582733
n02877765 bottlecap	0.10574431717395782
n04116512 rubber eraser, rubber, pencil eraser	0.0716797411441803

Category	Score
n02105505 komondor	0.30265742540359497
n02085936 Maltese dog, Maltese terrier, Maltese	0.10685333609580994
n02098413 Lhasa, Lhasa apso	0.07794724404811859
n02113624 toy poodle	0.06806328147649765
n02105641 Old English sheepdog, bobtail	0.028750602155923843

Category	Score
n04037443 racer, race car, racing car	0.22242484986782074
n03977966 police van, police wagon, paddy wagon, patrol wagon, wagon, black Maria	0.10872112214565277
n03417042 garbage truck, dustcart	0.0702013298869133
n03602883 joystick	0.06894396990537643
n03791053 motor scooter, scooter	0.06231105327606201

Category	Score
n03832673 notebook, notebook computer	0.39376965165138245
n03085013 computer keyboard, keypad	0.1367684155702591
n03642806 laptop, laptop computer	0.13363485038280487
n04264628 space bar	0.07292277365922928
n03485407 hand-held computer, hand-held microcomputer	0.054601266980171204

Category	Score
n02134084 ice bear, polar bear, Ursus Maritimus, Thalarctos maritimus	0.9999891519546509
n02120079 Arctic fox, white fox, Alopex lagopus	4.301877197576687e-06
n02132136 brown bear, bruin, Ursus arctos	2.4333951387234265e-06
n02114548 white wolf, Arctic wolf, Canis lupus tundrarum	1.4043683904674253e-06
n02104029 kuvasz	1.132186412178271e-06

Category	Score
n02110627 affenpinscher, monkey pinscher, monkey dog	0.47548773884773254
n02096177 cairn, cairn terrier	0.1800858974456787
n02328150 Angora, Angora rabbit	0.07240748405456543
n02097298 Scotch terrier, Scottish terrier, Scottie	0.06806369125843048
n02096294 Australian terrier	0.02631850726902485

Category	Score
n02206856 bee	0.30947500467300415
n02219486 ant, emmet, pismire	0.2313130646944046
n11939491 daisy	0.10648820549249649
n03476684 hair slide	0.10484256595373154
n03991062 pot, flowerpot	0.05003173649311066