Jetson Nanoでonnx-chainerを使ってONNX形式で出力

はじめに

TensorRTを使うためにchainerのモデルをONNX形式で出力することを考えた。どうせVGG16等有名どころのモデルを使うんだから最初からONNX形式のモデルをダウンロードするのが正解と思いつつ。

事前準備

TensorRTを使うのは初めてなので、今使っている環境を誤って壊さないよう、新たに$HOME/Documents/TensorRTに環境を作ることにした。ここにONNXの環境を構築する。

ターミナル

yamamo-to@jetson-nano:~$ mkdir ~/Documents/TensorRT
yamamo-to@jetson-nano:~$ cd ~/Documents/TensorRT
yamamo-to@jetson-nano:~/Documents/TensorRT$ pipenv --three
yamamo-to@jetson-nano:~/Documents/TensorRT$ pipenv install cupy
yamamo-to@jetson-nano:~/Documents/TensorRT$ pipenv install chainer
yamamo-to@jetson-nano:~/Documents/TensorRT$ pipenv install matplotlib
yamamo-to@jetson-nano:~/Documents/TensorRT$ pipenv install pillow

なおpipenv installでパッケージをインストールする際、一行ずつに分けているのは下記のようなエラーが多発するから。

pipenv.patched.notpip._vendor.urllib3.exceptions.ProtocolError: ("Connection broken: ConnectionResetError(104, 'Connection reset by peer')", ConnectionResetError(104, 'Connection reset by peer'))

これが発生したら同じコマンドを何回か粘るとうまく行く。うまく行くまで繰り返すのがコツ。

ONNX関連パッケージのインストール

何度やってもpipenv install onnxが失敗するのでソースコードを持ってきてインストールをしようとしたらパッケージがいくつか足りないことが分かった。

ターミナル

sudo apt install libprotobuf-dev protobuf-compiler pybind11-dev

ターミナル

yamamo-to@jetson-nano:~/Documents/TensorRT$ pipenv install pybind11
yamamo-to@jetson-nano:~/Documents/TensorRT$ pipenv install onnx
yamamo-to@jetson-nano:~/Documents/TensorRT$ pipenv install onnx-chainer
yamamo-to@jetson-nano:~/Documents/TensorRT$ pipenv install chainercv

現時点でのPipfileは下記の通り。

[[source]]
name = "pypi"
url = "https://pypi.org/simple"
verify_ssl = true

[dev-packages]

[packages]
cupy = "*"
chainer = "*"
matplotlib = "*"
pillow = "*"
pybind11 = "*"
onnx = "*"
onnx-chainer = "*"
chainercv = "*"

[requires]
python_version = "3.6"

実行

onnx-chainerのgithub¹のソースコードを流用した。ソースコードを見て、caffemodel形式から長い長い変換が始まるのではと思ったが、本ソースはnpzを直接ダウンロードしたので安心した。よくよく考えるとモデルの変換だけならJetson Nanoでやる必要もないような・・・。

chainer2onnx.py

import numpy as np

import chainer
import chainercv.links as C
import onnx_chainer

model = C.VGG16(pretrained_model='imagenet')

# Pseudo input
x = np.zeros((1, 3, 224, 224), dtype=np.float32)

onnx_chainer.export(model, x, filename='vgg16.onnx')

とりあえず実行。

ターミナル

yamamo-to@jetson-nano:~/Documents/TensorRT$ pipenv run python3 chainer2onnx.py

おぉ、動いた！インストールに手間取ったのとハングアップを繰り返した²ので、正直動くとは思っていなかった。実行時間もダウンロードを除けばJetson Nano上で5分程度。十分使える速度である。

vgg16.onnxは作成された模様。しかし謎のメッセージが・・・。これ、本当に使えるんだろうか。

ターミナル

[libprotobuf WARNING google/protobuf/io/coded_stream.cc:604] Reading dangerously large protocol message.  If the message turns out to be larger than 2147483647 bytes, parsing will be halted for security reasons.  To increase the limit (or to disable these warnings), see CodedInputStream::SetTotalBytesLimit() in google/protobuf/io/coded_stream.h.
[libprotobuf WARNING google/protobuf/io/coded_stream.cc:81] The total number of bytes read was 553437453

1. https://github.com/chainer/onnx-chainer ↩

2. 変換中に突然通信が停止しJetson Nanoごと落ちた ↩