Jetson NanoでVGG16を試す

はじめに

Jetson Nanoで機械学習の「学習側」をやるという試み¹は無理すぎた(やる前から分かっていたことだが・・・)。と言う訳で「推論側」に手を染めようと思う。MNISTと同じくらい有名なVGG16の学習済みモデルを使って画像を判別したいと思う。なんだかようやくスタートラインにたった気分。

画像の準備

画像はカメラがあるのでOpenCVを使って適当にtake_picture.pyを書いた。Raspberry Piでも感じたが「こんな感じで動くだろ」って思って作るとほぼ動くのが凄い。

take_picture.py

import cv2
import argparse


def main():
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument('filename')
    parser.add_argument('-w', '--width', default=600, type=int)
    parser.add_argument('-h', '--height', default=400, type=int)
    args = parser.parse_args()

    cap = cv2.VideoCapture(0)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, args.width)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, args.height)
    r, frame = cap.read()
    cv2.imwrite(args.filename, frame)


if __name__ == "__main__":
    main()

これを使った画像の取得は下記のようになる。

ターミナル

python3 take_picture.py -w 640 -h 360 boss.jpg

学習済みモデルの準備

今回はVGG16を使うので、chainerで

from chainer.links import VGG16Layers
vgg16 = L.VGG16Layers()

のような感じで書くと、モデルをダウンロードし始める。528MiBあるのでダウンロードにはそこそこ時間がかかる。さらにこのフォーマットはcaffemodel形式なので、これをchainerの形式に変換するプロセスが走るがこちらが全く終わらない。別のPCでやるか、もしくは寝る前くらいに下記のようなプログラムを使って変換したほうがいい。

caffe2chainer.py

import chainer.links as L
import argparse


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('input_caffe_model')
    parser.add_argument('output_chainer_model')
    args = parser.parse_args()

    L.VGG16Layers.convert_caffemodel_to_npz(
        args.input_caffe_model, args.output_chainer_model)


if __name__ == "__main__":
    main()

実行は次の通り。まずモデルをダウンロード。

ターミナル

wget http://www.robots.ox.ac.uk/~vgg/software/very_deep/caffe/VGG_ILSVRC_16_layers.caffemodel

次に変換。この変換はとても時間がかかるので別のPCでやった方がかなり速い。

ターミナル

python caffe2chainer.py VGG_ILSVRC_16_layers.caffmodel VGG_ILSVRC_16_layers.npz 

作ったモデルは

ターミナル

yamamo-to@jetson-nano:~/.chainer/dataset/pfnet/chainer/models/

に置いた。

VGG16による推論

チューニング等は無しで素のVGG16を使う。画像はさっき撮ったばかりのCRAFT BOSS。

ソースコードは下記の通り²。

vgg16_test.py

import cupy as cp
from chainer import functions as F
from chainer import links as L
from PIL import Image


vgg16 = L.VGG16Layers()
vgg16.to_gpu()

img = Image.open('./boss.jpg')
img = L.model.vision.vgg.prepare(img)
img = img[cp.newaxis]
img = cp.array(img, dtype=cp.float32)

result = vgg16(img)
val = F.argmax(result['prob'], axis=1)

print(result)
print(val)

実行は下記の通り。

ターミナル

pipenv run python3 vgg16_test.py

結果

実行時間は23秒弱。ほとんどが学習済みモデルの読み込みと思われる。

推論の結果は「variable([789])」と出た。

ラベルは下記のURLの行番号790行目(789は0スタートなので1を足した)を見ることにする。
https://github.com/davidgengenbach/vgg-caffe/blob/master/data/labels.txt

n04192698 shield, buckler
n04200800 shoe shop, shoe-shop, shoe store
n04201297 shoji                            ← これ
n04204238 shopping basket
n04204347 shopping cart

推論結果はなんと「障子」でした。
うん、間違ってない・・・。

1. Jetson NanoでChainerのDCGANサンプルを動かす試み ↩

2. F.softmaxは不要だったため削除 ↩