🤖

RubyでVGG Networkを使った画像認識

2021/02/11に公開

Ruby

はじめに

Rubyではlibtorch (PyTorchのC++版) のbindingライブラリであるtorch.rbというものがあります。これを使うと、Rubyでも深層学習による画像認識が行えます。

https://github.com/ankane/torch.rb

インストール

この記事での作業は、macOS Big Sur 11.2で行いました。必要なライブラリはhomebrewでインストールしていきます。まずは、libtorchとtorch.rbをインストールします。定義済みの画像系モデルが利用できるtorchvisionもインストールします。

$ brew install libtorch automake
$ gem install torch-rb torchvision

次に、画像をNumo::NArray形式で読み込むことができるmagroをインストールします。Numo::NArrayは、Pythonでいうところのnumpyで、torch.rbをはじめ多くのRubyの機械学習ライブラリが、Numo::NArrayに対応しています。magroはlibpngとlibjpegを利用するので、それらもインストールします。

$ brew install libpng jpeg
$ gem install magro

そして、認識結果を出力するさいに必要になる、ImageNetのクラス情報のJSONをダウンロードしておきます。

$ wget https://s3.amazonaws.com/deep-learning-models/image-models/imagenet_class_index.json

したごしらえ

実はlibtorchにバグがあって、PyTorchで保存されたモデルがそのままでは読み込めません（2020/2/11現在）。そのため、torchvisionで提供されている、学習済みのVGG Networkが利用できません。ただし、ちょっと修正してあげると、読み込めるようになります。

まずは、pytorchとtorchvisionをインストールします。

$ pip install torch torchvision

そして、以下のスクリプトで、学習済みモデルをlibtorchでも読める形式に変換します。

convert.py

import torch
import torchvision.models as models

def save_model(model, filename):
  model_dict = model.state_dict()
  model_dict = { k: v.data if isinstance(v, torch.nn.Parameter) else v for k, v in model_dict.items() }
  torch.save(model_dict, filename)

def main():
  save_model(models.vgg16(pretrained=True), 'vgg16_.pth')

if __name__ == '__main__':
  main()

これを、実行すると、vgg16_.pthというファイルができます。

$ python convert.py

画像認識してみよう

学習済みのVGG Networkを使って、画像認識を行うスクリプトを、パートにわけて解説します。こちらのWikipediaのゴールデンレトリバーを試してみましょう。ImageNetのクラスは意外と細かくて、golden_retrieverクラスがあります。

$ wget https://upload.wikimedia.org/wikipedia/commons/7/74/A_Golden_Retriever-9_%28Barras%29.JPG

学習済みモデルの読み込み

前節で得られた、学習済みモデルを読み込みましょう。ネットワークの定義はtorchvisionのものを使えばよいので、簡単です。

classify.rb

require 'torch'
require 'torchvision'

vgg = TorchVision::Models::VGG16.new
vgg.load_state_dict(Torch.load('vgg16_.pth'))

画像の読み込みと前処理

画像の読み込みはmagroを用います。magroは画像をNumo::NArray形式で読み込みます。Numo::NArrayはPythonでいうところのnumpyで、torch.rbでもtensorとの相互変換がサポートされています。torchvisionが提供する学習済みモデルの前処理は、以下に詳しく書かれています。

これに従った前処理を施していきます。

classify.rb

require 'magro'

# 画像を読み込む. 
img = Magro::IO.imread('A_Golden_Retriever-9_(Barras).JPG')

# 画像の中心を正方形に切り出す.
height, width, = img.shape
img_size = [height, width].min
y_offset = (height - img_size) / 2
x_offset = (width - img_size) / 2
img = img[y_offset...(y_offset + img_size), x_offset...(x_offset + img_size), true]

# 画像を224x224の大きさにする.
img = Magro::Transform.resize(img, height: 224, width: 224)

# 画素値を[0, 1]の範囲に正規化する.
img = Numo::SFloat.cast(img) / 255.0

# 画像をtorch.rbのtensorに変換し, [チャンネル, 高さ, 幅]の順に入れ替える.
img_torch = Torch.from_numo(img).permute(2, 0, 1)

# 平均と標準偏差を正規化する.
mean = Torch.tensor([0.485, 0.456, 0.406])
std = Torch.tensor([0.229, 0.224, 0.225])
normalize = TorchVision::Transforms::Normalize.new(mean, std)
normalize.call(img_torch)

# tensorを [1, 3, 224, 224] の形にする.
img_torch = img_torch.expand(1, -1, -1, -1)

画像認識と結果の出力

前処理した画像を、学習済みモデルに与えて、画像認識してみましょう。最終層の出力の各要素は、ImageNetのクラス番号に対応しています。最も大きな値となった要素が、認識結果となります。

classify.rb

require 'json'

# 学習済みモデルに、前処理した画像を入力する.
vgg.eval
out = vgg.forward(img_torch)

# 最終層の出力で最も値の大きい要素の添字を得る.
class_idx = out.numo[0, true].max_index

# 添字に対応するImageNetのクラスを出力する.
imagenet_classes = JSON.load(File.read('imagenet_class_index.json'))
puts "class: #{imagenet_classes[class_idx.to_s].last}"

以上ちょっと長めのコードになりました。一つにまとめたものをGistに置きました。

これを実行すると...ゴールデンレトリバーと認識されました〜🎉

$ ruby classify.rb
class: golden_retriever

おわりに

torch.rbを使うことで、Rubyでも学習済みのVGG Networkによる画像認識ができました。実際には、ImageNetの1,000クラスに画像を分類したいという場面は少ないように思います。最終層の手前の出力を特徴ベクトルとして、自前で用意したラベルで学習したり、画像検索をすることが多いのではないでしょうか。その方法は、また別の記事で紹介できたらと思います。Rubyで機械学習・深層学習ができる時代がきています。

はじめに

インストール

したごしらえ

画像認識してみよう

学習済みモデルの読み込み

画像の読み込みと前処理

画像認識と結果の出力

おわりに

Discussion