TensorflowのInception-v3実装のCPUボトルネックをGPUに移す

この記事は

• ディープラーニングの画像認識の学習時にCPUがボトルネックになってGPUが本領を発揮できない事象があったのでそれを解消した時のメモです

事象

• Tensorflowのmodelリポジトリには画像認識CNNのInception-v3の実装が入っています。
• パラメータ1つで複数GPUでの学習を実行できるようになっています
• しかしなぜか複数GPUで動作させても学習速度が上がりません、てかむしろちょっと遅くなってる？
  • 1batchあたり0.7秒

改善前：GPU*1の場合

2016-12-16 03:35:23.690928: step 120, loss = 13.14 (45.2 examples/sec; 0.707 sec/batch)
2016-12-16 03:35:31.121967: step 130, loss = 13.06 (44.6 examples/sec; 0.717 sec/batch)
2016-12-16 03:35:38.888079: step 140, loss = 12.70 (45.1 examples/sec; 0.709 sec/batch)

改善前：GPU*2の場合

2016-12-16 03:34:36.980520: step 60, loss = 13.75 (41.7 examples/sec; 0.767 sec/batch)
2016-12-16 03:34:44.441697: step 70, loss = 13.13 (42.3 examples/sec; 0.757 sec/batch)
2016-12-16 03:34:52.268679: step 80, loss = 13.02 (38.8 examples/sec; 0.824 sec/batch)

• せっかく高い金払ってTitan Xを2基積んだのにこれじゃ何の意味もないじゃないか・・・

原因

• vmstatとかtopでリソースを確認するとCPU使用率が非常に高くなっていました
  • 一方GPU使用率は100%と0%を行ったり来たりしてます
  • I/Oとかメモリとかは全然普通です
• どうもCPUがボトルネックになってGPUを活かしきれていないっぽいです

調査

• どこでCPU使ってるんだろう、と思ってinception-v3のコードを眺めていたらこんなコードを見つけました

image_processing.py

def distorted_inputs(dataset, batch_size=None, num_preprocess_threads=None):
(snip)
  # Force all input processing onto CPU in order to reserve the GPU for
  # the forward inference and back-propagation.
  with tf.device('/cpu:0'):
    images, labels = batch_inputs(
        dataset, batch_size, train=True,
        num_preprocess_threads=num_preprocess_threads,
        num_readers=FLAGS.num_readers)
  return images, labels

• なんかcpuって思いっきりハードコードされてるけど・・・！

処理内容

• distorted_inputsというメソッドはその名の通り入力画像にディストーションをかける処理です。
• ディストーションとはCNNの前処理で、元画像を拡大・縮小・色み変更などなどでランダムに加工し、入力画像を水増しして精度向上に役立てる処理のことです。
• ちなみにTensorboardで学習中の様子を見ると下記のようなディストーションが効いた画像が見られて、面白いです

• コメントに入力データの加工はあくまで前処理であり、ディープラーニングの本領である学習部分にGPUパワーを使うために、あえてCPUを用いることを強制させているみたいなことが書いてある気がします。

対処

• 普通にソースをgpuに書き換えてbazel buildしなおしました

image_processing.py

  with tf.device('/gpu:0'):

再実行

• なんとGPU1基ですら性能が改善しています！(約0.45秒/batch)
• GPU2基だと倍にはならないけど速くはなってる！(約0.3秒/batch)

改善後：GPU*1の場合

2016-12-21 12:20:44.338667: step 20, loss = 14.27 (69.2 examples/sec; 0.462 sec/batch)
2016-12-21 12:20:48.906566: step 30, loss = 14.48 (69.4 examples/sec; 0.461 sec/batch)
2016-12-21 12:20:53.492054: step 40, loss = 13.77 (69.6 examples/sec; 0.459 sec/batch)

改善後：GPU*2の場合

2016-12-21 12:16:46.828951: step 10, loss = 13.05 (108.9 examples/sec; 0.294 sec/batch)
2016-12-21 12:16:49.877032: step 20, loss = 14.37 (95.8 examples/sec; 0.334 sec/batch)
2016-12-21 12:16:53.138938: step 30, loss = 14.94 (88.7 examples/sec; 0.361 sec/batch)

終わりに

• なんとかGPU2枚差しの面目を保ちました
  • 2倍にはなってないので、またCPUがボトルネックになってそうだけど・・・
• Googleの人たちのマシンはCPUが超リッチでボトルネックにならないんですかね
• いずれにしてもハードコードはどうかなと思わなくもない