Tensorflow's while loop slow on GPU? - tensorflow

Question

For unknown reasons, the following code takes twice much slower on GPU then on CPU. Could anyone explain why:
import time
import tensorflow as tf
with tf.device('/device:GPU:0'): # gpu takes: 5.132448434829712 seconds
# with tf.device('/cpu:0'): # cpu takes: 3.440524101257324 seconds
i = tf.constant(0)
while_condition = lambda i: tf.less(i, 2 ** 20)
a = tf.fill([16, 16], 1.1)
b = tf.fill([16, 16], 2.2)
def body(i):
res = tf.matmul(a, b)
# increment i
add = tf.add(i, 1)
return (add,)
ini_matmul = tf.matmul(a, b)
# do the loop:
loop = tf.while_loop(while_condition, body, [i])
with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as sess:
sess.run(ini_matmul) # force GPU to initilise anything it needs.
t0 = time.time()
sess.run(loop)
t1 = time.time()
print(t1 - t0)
sess.close()
Note: Usually, the GPU runs for 5 seconds, CPU runs for 3 seconds, and CPU version using numpy runs for only 1.5 seconds.
Hardware: Tensorflow code running on Google's Colab. Numpy code running on local Intel Core i5-7267U.
Numpy version:
import numpy as np
import time
i = 0
a = np.full([16,16],1.1)
b = np.full([16,16],2.2)
t0 = time.time()
while i < 2**20:
a.dot(b)
i += 1
t1 = time.time()
print(t1-t0)
Update
This is becoming more and more wired to me because scaling up the matrix does not really helps. Here's the updated code and data in it (running of Titan XP card/Intel i7 CPU). Essentially tensorflow is running much slower.
import time
import tensorflow as tf
dimension = 11
repeat = 2**10
use_gpu = False
# Device: /device:GPU:0, Dimension 11, Repeat: 1024, Time cost: 0.00457597 seconds.
# Device: /cpu:0, Dimension 11, Repeat: 1024, Time cost: 0.00353599 seconds.
dev_name = '/device:GPU:0' if use_gpu else '/cpu:0'
with tf.device(dev_name):
i = tf.constant(0)
while_condition = lambda i: tf.less(i, repeat)
a = tf.constant(1.1, shape=[2**dimension, 2**dimension])
b = tf.constant(2.2, shape=[2**dimension, 2**dimension])
def body(i):
res = tf.matmul(a, b)
add = tf.add(i, 1)
return (add,)
ini_matmul = tf.matmul(a, b)
# do the loop:
loop = tf.while_loop(while_condition, body, [i])
with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as sess:
sess.run(ini_matmul) # force initialisation.
t0 = time.time()
sess.run(loop)
t1 = time.time()
print('Device: {dev}, Dimension {dim:d}, Repeat: {r:d}, Time cost: {t:.8f} seconds.'.format(
dev = dev_name,
dim = dimension, r = repeat,
t = t1 - t0
))
sess.close()

Answer 1

In the end, I figured out that the matmul operation is not executed by tensorflow, because it is an orphan node in the graph.

Answer 2

This is an interesting question.
The relative slowdown you're seeing between GPU and CPU execution in the TensorFlow snippet is almost certainly due to GPU memory allocation overhead. To summarize the link, cudaMalloc is slower than malloc. This memory allocation slowdown is offset by a speedup in the requested operation (matmul in this case) if and only if the speedup exceeds the difference in memory allocation times. This is always true for matmul when the matrices are large. It is not true when the matrices are small, as is the case in your example. To validate this hypothesis, iteratively increase the size of the multiplicands and record both the CPU and GPU running times - the two should converge, then cross, if memory allocation is indeed the issue.
The difference between the Numpy running time and the CPU-only running time is likely due to very subtle difference between the Numpy and TensorFlow code. Note that in the Numpy code you only instantiate a and b once. It looks like you do the same thing in the TensorFlow code because you only call your initialization once, but you're still populating the tensors in every iteration! To see why, note that tf.fill returns a Tensor. By definition, Tensor objects are populated each time sess.run is called on the graph that contains them. Thus, the two snippets actually do slightly different things. A more direct comparison would be to make a and b both tf.constant in the TensorFlow snippet.