TensorFlowで複数GPUで2次元畳み込みやってみる(tf.nn.conv2d with multiple GPUs)

1. やったこと

TensorFlow の tf.nn.conv2d関数 を GCE 上でK80とV100のGPUを1個〜8個を用いて、データ並列で実行。処理速度の検証を行った。

2. わかったこと

以下のコードを試した結果、CPUのほうが3倍早く、期待はずれの結果となってしまった。こうすれば早くなるなどのコメントがありましたらぜひご教授願います。

3. PCのスペック

CPU: virtual 12 core
MEM: 32 GB
GPU: K80x8個、V100x8個で検証

4. 検証に用いたコードの概要

multi_gpu.py

GPU の数と画像データ数、画像のサイズを引数にとり、データ並列でGPUで2次元畳み込みを行い、処理時間を出力するスクリプト

run_multi_gpu.py

multi_gpu.py を GPU の数を変えて実行し結果を記録するスクリプト

5. 実際のコード

TensorFlow の 複数 GPU の使い方は公式サイトの以下のページを参考にしました。
"https://www.tensorflow.org/guide/using_gpu#using_multiple_gpus"

multi_gpu.py
import numpy as np
import tensorflow as tf
import time
import sys

N_gpus = int(sys.argv[1])
N_images = int(sys.argv[2])
width = int(sys.argv[3])
height = int(sys.argv[4])

c = []

samples = np.ones((N_images,width,height,1), dtype=np.float32)
split_samples = tf.split(samples, N_gpus, axis=0)

for n in range(N_gpus):
    d = '/gpu:%d'%n
    with tf.device(d):
        fltr = tf.constant(np.ones((3,3,1,1), dtype=np.float32))
        conv = tf.nn.conv2d(split_samples[n], fltr, (1,1,1,1), 'SAME')
    c.append(conv)
    
with tf.device('/cpu:0'):
  sum = tf.add_n(c)
  
# Creates a session with log_device_placement set to True.
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
# Runs the op.
s = time.time()
out = sess.run(sum)
e = time.time()
print(out)
print((e-s)*1000, 'ms')


out_text = '================================================\n'
out_text = \
    'N_gpus = {}, N_images = {}, width = {}, height = {}\n'\
    .format(N_gpus, N_images, width, height)
time_estimated = '%f ms\n'%((e-s)*1000)
out_text += time_estimated

filename = 'multi_gpu_time_estimated_{}_{}_{}_{}.txt'\
    .format(N_gpus, N_images, width, height)

with open(filename, 'w') as fp:
    fp.write(out_text)
run_multi_gpu.py
# -*- coding:utf-8 -*-
import subprocess

N_GPUS_MAX = 8
python_exe = 'python3'
script_name = 'multi_gpu.py'

# データ数は GPU の個数の倍数でなければならない
# 1, 2, 3, 4, 5, 6, 7, 8 の倍数は840
# 840個の画像データ、サイズ(100,100)の
# を2次元畳み込みすることを考える

n_images = 840
width = 100
height = 100

for n_gpus in range(1, N_GPUS_MAX+1):
    command = [
    	python_exe,
    	script_name,
    	str(n_gpus),
    	str(n_images),
    	str(width),
    	str(height)]
    out = subprocess.check_output(command)
    filename = 'multi_gpus_time_log_{}_{}_{}_{}.txt'\
        .format(n_gpus, n_images, width, height)
    with open(filename, 'wb') as fp:
        fp.write(out)

6. 実行結果

以下にV100 という GPU で実行した結果を載せます。
上から順に GPU1個、GPU2個、、、GPU8個と増やして実行た結果です。

================ GPU V100 ================
N_gpus = 1, N_images = 840, width = 100, height = 100
14897.803307 ms

N_gpus = 2, N_images = 840, width = 100, height = 100
4328.074455 ms

N_gpus = 3, N_images = 840, width = 100, height = 100
5645.370483 ms

N_gpus = 4, N_images = 840, width = 100, height = 100
684.230804 ms

N_gpus = 5, N_images = 840, width = 100, height = 100
744.736195 ms

N_gpus = 6, N_images = 840, width = 100, height = 100
731.468678 ms

N_gpus = 7, N_images = 840, width = 100, height = 100
726.519346 ms

N_gpus = 8, N_images = 840, width = 100, height = 100
721.646070 ms

同様に、K80 という GPU で実行した結果を以下に載せます。

================ GPU K80 ================
N_gpus = 1, N_images = 840, width = 100, height = 100
3248.390198 ms

N_gpus = 2, N_images = 840, width = 100, height = 100
3380.734682 ms

N_gpus = 3, N_images = 840, width = 100, height = 100
3837.318182 ms

N_gpus = 4, N_images = 840, width = 100, height = 100
690.669060 ms

N_gpus = 5, N_images = 840, width = 100, height = 100
691.730976 ms

N_gpus = 6, N_images = 840, width = 100, height = 100
588.120937 ms

N_gpus = 7, N_images = 840, width = 100, height = 100
581.115961 ms

N_gpus = 8, N_images = 840, width = 100, height = 100
647.247314 ms

7. GPU がなんとなく遅い気がするのでCPUでやってみた結果

import numpy as np
import scipy
from timeit import timeit

fltr = np.ones((3,3), dtype=np.float32)
a = np.ones((100,100), dtype=np.float32)

N_images = 840

と定義して以下の scipy.signal.convolve2d で畳み込みをやってみると、、、

>>> %timeit [scipy.signal.convolve2d(a, fltr, 'same') for _ in range(N_images)]
222 ms ± 4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

なぜかCPU1個の方がGPU8個より3倍くらい早いです。

8. GPU の使用数が 1-3 のときに遅い原因について推測してみる

GPU が 2-3 個のとき、処理時間が 1000 ms を超えており極端に遅いです。
極端に遅いケースについては、データが GPU のメモリに入りきらず、
外部とデータのやりとりを行うスワップのようなものが発生していることが考えられます。

9. CPU より GPU が遅い原因について推測、検証方法を検討してみる(検証はしません)

・TensorFlow がたいして最適されていない
⇨ Keras でも同様の結果になるか確認
⇨ Chainer など他のライブラリを使用してみる
・畳み込みのフィルタサイズの影響
⇨ フィルタサイズ 3x3 でなく10x10 などサイズを大きくする
・チャンネルのサイズの影響
⇨ チャンネルのサイズを 1 でなく 10 などサイズを大きくする
GPUからCPUに計算結果を転送するのがボトルネックになっている
⇨ データ転送するだけのコードで検証する
GPU内でずっと計算し続けて、オーバーヘッドの影響をなくすようなコードで検証する

10. まとめ

勝手な想像ですが、 TensorFlow は tf.nn.conv2d のためにあるわけではなく、どちらかというと、ニューラルネットの重み計算を行うために最適化されているはずです。そのため、静的な数式で表されるような重み更新の計算式などはうまくGPUで分散して計算できると思うのですが、畳み込みのようなシーケンシャルな処理は単純な計算式で表現できないため、それも遅さの原因の一つかなと思います。ただ、そうはいっても今回の例ではデータ並列したはずなので、GPUの数に比例して早くなるはずだったのですが、そうはなりませんでした。現実は厳しかったです。このような畳み込みだけを行うという特殊なケースでは単純にGPUを増やせば早くなるというわけではないようですので注意が必要かと思われます。
いろいろ書きましたが、検証は面倒なので、TensorFlowの開発者が最適化をすすめて、コードを書き換えなくても自動で計算リソースをスケールしてくれる日がいつかくることを期待して、メインの機械学習アルゴリズムについて学習をすすめて行こうと思います。

以上です。おやすみなさい。