PyTorch
caffe2
ONNX

onnxを使ってみました.

新年1つ目の投稿です.

facebookではモデルの学習をPyTorchで行い,プロダクションではcaffe2を使っており1
onnxは上2つをはじめとするニューラルネットワークフレームワーク間でモデルを渡す上で非常に有用です.

そこで今回はとりあえずPyTorchで学習したモデルをonnxでファイルに出力し,caffe2から利用し,
計算速度を比較してみたところ,PyTorchの方が速かった,という話です.
なお,とりあえず動かすことが目標だったのでonnxとcaffe2のドキュメントは斜め読みですので誤りが含まれる可能性が高いです.

caffe2環境はdockerで構築し,モデルはPyTorch/examplesにあるmnistを基に作成しました.
用いたスクリプトは以下です.
- Dockerfile
- main.py

速度比較に用いた入力は

  • バッチサイズ: $64$
  • チャネル数: $1$
  • 縦: $28$
  • 横: $28$

です.