python - Random number generation in Tensorflow is very slow (tensorflow-gpu)

Question

I have been starting to learn Tensorflow (on GPU) and have been playing with some of the basic functions, including random number generation with Tensorflow Probability. I have found that the random number generation to be surprisingly slow. Here is the code I used to test:

import time
import tensorflow as tf
import numpy as np
import tensorflow_probability as tfp
tfd = tfp.distributions

size = (100,100)
counts = np.random.randint(1000,size = size)
probs = 0.5*np.ones(size)

t_counts = tf.cast(tf.convert_to_tensor(counts),tf.float64)
t_probs = tf.convert_to_tensor(probs)

N = 20
samples = np.zeros((size[0],size[1],N))
tic = time.time()
for i in range(N):
    samples[:,:,i] = np.random.binomial(counts,probs)
print('Numpy took',time.time()-tic,'seconds')

t_samples = []
b = tfd.Binomial(t_counts,t_probs)
tic = time.time()
for i in range(N):
    t_samples.append(b.sample())
print('Tensorflow took',time.time()-tic,'seconds')

I expected a marginal speedup over numpy at these array sizes but instead found Tensorflow to be dramatically slower. For the given size the output I get is:

Numpy took 0.04088950157165527 seconds
Tensorflow took 404.8193895816803 seconds

which seems dramatic. I tried scaling up the size of the size of the array to see how it scaled, and for size = (200,200) I get:

Numpy took 0.16755127906799316 seconds
Tensorflow took 1649.4787373542786 seconds

So tensorflow (ostensibly using the GPU) seems to be scaling in exactly the same way as numpy using the CPU. Is this expected behavior? I haven't been able to find any benchmarks for Tensorflow random number generation but it's not something I have specialized expertise in. Is there something that I am doing wrong when setting this up?

标签： pythonnumpytensorflowrandomgpu