TensorFlowはCaffeよりも早いのか?(後編)

tensorflow
前回の記事において、TensorFlow 0.7.1 (GPU版)でconvnet-benchmarksを実行するとエラーが発生しましたので、この対策を行って、CaffeとTensorFlowのどちらが早いのか決着をつけたいと思います。

スポンサーリンク

TensorFlow 0.8.0のインストール

一番可能性がある対策は、TensorFlowの最新版でベンチマークを実行することなので、一旦TensorFlow 0.7.1 (GPU版)をdeactivateし、TensorFlow 0.8.0 (GPU版)を新規にインストールします。

(tensorflow-gpu)$ deactivate
$ cd
$ virtualenv ~/tensorflow-0.8.0
$ source ~/tensorflow-0.8.0/bin/activate
(tensorflow-0.8.0)$ pip install --upgrade https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow-0.8.0-cp27-none-linux_x86_64.whl

また、TensorFlow 0.7.1のときと同じ動作確認を行いました。

TensorFlow 0.8.0ベンチマーク

ベンチマークを実行するために、ディレクトリを移動します。

(tensorflow-0.8.0)$ cd ~/convnet-benchmarks/tensorflow

pythonスクリプトbenchmark_googlenet.pyは、引数でバッチサイズを変更できるので、前回の記事と同様に、--batch_size 32を追加して、GoogLeNetのバッチサイズをデフォルトの128から32に変更します。
早速ベンチマークを実行します。

(tensorflow-0.8.0)$ python benchmark_googlenet.py --batch_size 32
...
Forward across 100 steps, 0.099 +/- 0.000 sec / batch
...
Forward-backward across 100 steps, 0.323 +/- 0.001 sec / batch

TensorFlow 0.8.0では、正常終了しました。
ベンチマークの結果としては、Caffeに迫る結果となっています。

benchmark_googlenet.pyの変更履歴を確認すると、2016年3月13日にNCHWフォーマットのサポートが追加されています。具体的には、引数としてdata_formatが追加され、対応するtf.nn.conv2d周りも変更されていることが分かりました。
従来と同じデータフォーマットにするため--data_format 'NHWC'を追加して、ベンチマークを実行します。

(tensorflow-0.8.0)$ python benchmark_googlenet.py --batch_size 32  --data_format 'NHWC'
...
Forward across 100 steps, 0.120 +/- 0.000 sec / batch
...
Forward-backward across 100 steps, 0.390 +/- 0.001 sec / batch

--data_format 'NHWC'を追加すると遅くなることから、benchmark_googlenet.pyのdata_formatがデフォルトの'NCHW'のときに、高速化を実現できていることが分かります。なお、関数tf.nn.conv2ddata_formatのデフォルトは'NHWC'です。

次に、NCHWフォーマットのサポートが追加される前、2015年11月13日にコミットされたbenchmark_googlenet.pyをダウンロードして、ファイル名をbenchmark_googlenet_old.pyに変更し、~/convnet-benchmarks/tensorflowにコピーします。
benchmark_googlenet_old.pyを用いて、ベンチマークを実行します。

(tensorflow-0.8.0)$ python benchmark_googlenet_old.py --batch_size 32
...
Forward across 100 steps, 0.156 +/- 0.001 sec / batch
...
Forward-backward across 100 steps, 0.468 +/- 0.005 sec / batch

benchmark_googlenet_old.pyでは、更に遅くなることから、現状のbenchmark_googlenet.pyは、NCHWフォーマットのサポート以外にも高速化が図られているようです。

TensorFlow 0.7.1ベンチマーク

最後にTensorFlowを0.8.0から0.7.1に切り替え、TensorFlow 0.7.1とbenchmark_googlenet_old.pyの組み合わせで、ベンチマークを実行します。

(tensorflow-0.8.0)$ deactivate
$ source ~/tensorflow-gpu/bin/activate
(tensorflow-gpu)$ python benchmark_googlenet_old.py --batch_size 32
...
Forward across 100 steps, 0.153 +/- 0.002 sec / batch
...
Forward-backward across 100 steps, 0.549 +/- 0.007 sec / batch

フォワードに関しては、TensorFlow 0.7.1の方が、ほんの少し早いのですが、フォワード・バックワード全体では、TensorFlow 0.8.0の方が、約80ms早くなっているので、TensorFlow 0.8.0は、バックワードに関して高速化が図られているようです。

ベンチマークのまとめ

と言うことで、GTX 960 (4GB)を用いた今回のGoogLeNetのベンチマーク結果をまとめると、以下の表のようになります。

GoogleNet V1 (GoogLeNet) – Input 32x3x224x224

Library Time (ms) forward (ms) backward (ms)
Caffe *1 321 98 222
TensorFlow 0.8.0 *2 323 99 224
TensorFlow 0.8.0 *3 390 120 270
TensorFlow 0.8.0 *4 468 156 312
TensorFlow 0.7.1 *4 549 153 396

*1: caffe-nv 0.14.5-2+cuda7.5
*2: 2016年4月15日にコミットされたbenchmark_googlenet.py(--data_format 'NCHW')
*3: 2016年4月15日にコミットされたbenchmark_googlenet.py(--data_format 'NHWC')
*4: 2015年11月13日にコミットされたbenchmark_googlenet.py

今回のベンチマークの結果としては、Caffe(cuDNN v5対応版のcaffe-nv)が、TensorFlow 0.8.0よりも僅かに早いのですが、TensorFlow 0.8.0がcuDNN v4対応版であることを考慮すると、cuDNN v5に対応したTensorFlow 0.8.0であれば、逆転もあるかもしれません。
また、TensorFlow 0.8.0に関しては、改良されている関数等を活用することで、0.7.1よりも大幅に処理時間を短縮できるようになっていることに驚きました。
GTX 960 (4GB)という低火力環境で、ImageNetモデルを学習させるのは諦めていますが、CIFAR-10でも、同等の効果が期待できるのか気になる所です。