Deep Learning フレームワークざっくり紹介 Advent Calendar 2017 の 9日目の記事です。

PyTorchとCaffe2で、モデル表現の標準フォーマットであるONNX (Open Neural Network Exchange)を使ってみます。

環境

今回はmacを使います。

macOS 10.13
python 2.7.14 ( anaconda2-5.0.1 )

PyTorch

元々はLuaという言語で書かれたTorchのPython版として、ChainerをForkしてFacebook AI Researchの方が開発したDeepLearningフレームワークです。

Chainerと同じく、Define by Runを採用しており動的ネットワークを書くことができます。書き方も基本的にはChainerとかなり似ています。

また、比較的新しいフレームワークで、2017年はじめにβ版が発表されて有名になりました。

Google Trendsでも2017年2月頃を堺にChainerを抜いています。

日本といくつかの国はChainerの検索数が多いですが、他の国では圧倒的にPyTrochの方が優位です。

インストール

公式ページに環境毎のインストール方法が書いています。condaやpipであれば、簡単にインストールできます。condaでのインストールを推奨しているみたいです。

condaの場合

$ conda install pytorch torchvision -c soumith

pipの場合

$ pip install install http://download.pytorch.org/whl/cu75/torch-0.2.0.post3-cp36-cp36m-manylinux1_x86_64.whl
$ pip install torchvision

しかし、現時点の最新バージョンである0.2.0版にはまだnn.onnxが含まれていないので、ONNXを利用するためには、masterブランチを使う必要があり、ソースからのインストールが必要です。

$ conda install numpy pyyaml setuptools cmake cffi
$ git clone --recursive https://github.com/pytorch/pytorch
$ cd pytorch
$ MACOSX_DEPLOYMENT_TARGET=10.13 CC=clang CXX=clang++ python setup.py install

10分以上かかりましたが、エラーなくインストールできました。

モデル定義

シンプルなCNNのモデルです。

import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x)

Caffe2

Caffe2は、元々Caffeを開発したFacebookのYangqing Jia氏が中心となり開発されました。

インストールは、公式サイトを参考に行います。ソースからのインストールを試みたのですが、上手くいかなかったので諦めてDockerを使いました。

$ docker pull caffe2ai/caffe2:c2v0.8.0.cpu.min.ubuntu16.04

ONNX

ONNX (Open Neural Network Exchange) は、FacebookとMicrosoftが提唱しているモデル表現の標準フォーマットです。異なるフレームワークで作成されたモデルを別のフレームワークでも使えるようにするといったものです。

PyTorchで書いたモデルをエクスポートしてCaffe2で使ったり、MXNetモデルをCognitive Toolkitで使ったりすることができます。

PyTorchの他、Caffe2, Microsoft Cognitive Toolkit, Apache MXNetがONNXをサポートしています。TensorflowやApple CoreMLでも使えるみたいです。

	Import	Export
Caffe2	OK	OK
Cognitive Toolkit	OK	OK
MXNet	OK	OK
PyTorch	NG?	OK

公式のチュートリアルもあるので、こちらも参考にしてみてください。

インストール

ONNXのインストールは、condaかpipで行います。condaでのインストールが推奨されているようです。

condaの場合

$ conda install -c conda-forge onnx

pipの場合

$ pip install onnx

モデルのエクスポート

先程のPyTorchモデルをエクスポートしてみます。

from torch.autograd import Variable
import torch.onnx

model = Net()
x = Variable(torch.randn(1, 1, 28, 28))
torch.onnx.export(model, x, 'model.onnx', verbose=True)

今回は、そのままエクスポートを行いますが、通常は学習済みモデルをエクスポートします。

PyTorchは動的にネットワークを定義(Define by Run)することができますが、ONNXのフォーマットに出力するためには、入力次元を指定する必要があります。

torch.onnx.exportで指定したファイル(model.onnx)にバイナリ形式で出力されます。verbose=Trueとすると人間が読める形式でコマンドラインに表示されます。

graph(%0 : Float(64, 1, 28, 28)
      %1 : Float(10, 1, 5, 5)
      %2 : Float(10)
      %3 : Float(20, 10, 5, 5)
      %4 : Float(20)
      %5 : Float(50, 320)
...
  %22 : Float(64, 50), %23 : UNKNOWN_TYPE = Dropout[is_test=1, ratio=0.5](%21)
  %26 : Float(64, 10) = Gemm[alpha=1, beta=1, broadcast=1, transB=1](%22, %7, %8)
  %27 : Float(64, 10) = Softmax[axis=1](%26)
  %28 : Float(64, 10) = Log(%27)
  return (%28);
}

モデルファイルの中身はonnxモジュールを使って確認することもできます。恐らくどのフレームワークで作成したモデルファイルでも同じように確認することができます。

import onnx
import onnx.helper

model = onnx.load("model.onnx")
print(onnx.helper.printable_graph(model.graph))

graph torch-jit-export (%name: "0"
type {
  tensor_type {
    elem_type: FLOAT
    shape {
      dim {
        dim_value: 1
      }
...

モデルのインポート

PyTorchではまだインポートがサポートされていないみたいなので、Caffe2で利用してみます。

Caffe2でONNXモデルを利用するためにonnx-caffe2をインストールします。

condaの場合

$ conda install -c ezyang onnx-caffe2

pipの場合

$ pip install onnx-caffe2

準備が整ったら、先程エクスポートしたmodel.onnxをインポートして利用してみます。

import numpy as np
import onnx
from onnx_caffe2.backend import prepare

model = onnx.load("model.onnx")
rep = prepare(model)
outputs = rep.run(np.random.randn(1, 1, 28, 28).astype(np.float32))
print outputs[0]
# >>> array([[-2.34730482, -2.28282928, -2.2903831 , -2.40077782, -2.38225985,
#             -2.35983896, -2.3191402 , -2.38344955, -2.22534752, -2.07859206]], dtype=float32)