Pytorch と Caffe2 のパフォーマンス比較

Resnet50 の pretrainedモデルを使って、CPUでの推論の性能差(速度)をチェックする。
caffe2, pytorch はともに以下のcondaパッケージを利用 (python3.6)

pytorch                   0.3.1           py36_cuda0.0_cudnn0.0_2
caffe2                    0.8.dev          py36hbf6519f_0 

テスト環境は3年くらい前のMBP

2.5GHz Intel Core i7
16GB メモリ

pytorchのコード

import time

import torch
from torch.autograd import Variable
import torchvision.models as models


model = models.resnet50(pretrained=True)
model.eval()
model.cpu()

dummy_input = Variable(torch.randn((1, 3, 224, 224))).cpu()

start = time.time()
for i in range(100):
    res = model(dummy_input)
print(time.time()-start)

caffe2のコード

import numpy as np
import time
from caffe2.python import workspace
from caffe2.python.models import resnet50

p = workspace.Predictor(resnet50.init_net, resnet50.predict_net)

dummy_input = np.random.random((1, 3, 224, 224)).astype(np.float32)

start = time.time()
for i in range(100):
    res = p.run({'gpu_0/data': dummy_input})
print(time.time()-start)

※入力のデータ名がgpu_0~ となってますがcaffe2のモデルの仕様の問題で, gpuは使ってません.

結果

それぞれ上記のコードで3回計測し、中央値で評価した。

	pytorch	caffe2
処理時間	21.153s	6.232s

cpuでの推論はcaffe2の方がpytorchより3倍~ほど速かった。
今回はcaffe2, pytorchともにcondaパッケージをそのまま使ったので自前で最適な構成でインストールすればまた違う結果になると思います。