CODEWITHSUNDEEP
pythonpython-3.xloopstensorflow
Jaffer Wilson
Jaffer Wilson

Reputation: 7273

How to convert the loop in Tensorflow loop?

I have written a program and trying to write it with Tensorflow so that I can use the GPU. Please let me know what I can do to convert it into tensorflow.
Here is the code for Normal Execution:

def function(array,position,times):
    o = []
    for loop in range(position,30):
        result = 0.0
        for k in range(times):
            result += array[loop-k]
        result /= times
        o.append(result)
    return o
q = [i for i in range(30)]
output = function(q,10,5)
print(output)

Using the Tensorflow I have tried this:

import tensorflow as tf
tf.enable_eager_execution()
def function_tf(array,position,times):
    o = tf.TensorArray(tf.float32,(30-position))
    index = tf.Variable(0)
    for loop in range(position,30):
        result = tf.Variable(0.0)
        for k in range(times):
            result.assign_add(array[loop-k].numpy())
            # print(result)
        result = tf.divide(result,times)
        o = o.write(index.numpy(),result.numpy())
        index.assign_add(1)
    return o.stack().numpy()
q = [i for i in range(30)]
q = tf.convert_to_tensor(q)

output = function_tf(q,tf.constant(10),tf.constant(5.0))
print(output)

The output is same: 

[ 8.  9. 10. 11. 12. 13. 14. 15. 16. 17. 18. 19. 20. 21. 22. 23. 24. 25.
 26. 27.]

But the only fact is the loops in the Tensorflow code still are the Python loop so it will work with CPU and not with GPU. I want everything to happen inside the GPU. Only assignment could be through the CPU.
Please help me convert the loop.

Upvotes: 1

Views: 566

Answers (1)

javidcf
javidcf

Reputation: 59681

You can actually write that function more simply without loops as the following convolution:

import tensorflow as tf
tf.enable_eager_execution()

def function(array, position, times):
    array = tf.convert_to_tensor(array)
    array = array[tf.newaxis, tf.maximum(position - times + 1, 0):, tf.newaxis]
    filter = tf.fill([times, 1, 1], tf.dtypes.cast(1 / times, array.dtype))
    return tf.nn.conv1d(array, filter, stride=1, padding='VALID')[0, :, 0]
print(function(tf.range(30.), 10, 5).numpy())
# [ 8.000001  9.       10.       11.       12.       13.       14.
#  15.       16.       17.       18.       19.       20.       21.
#  21.999998 23.       24.000002 25.       26.       26.999998]

EDIT: If you want to explicitly recreate the loop in the original code (although it will probably be slower), you can do it with tf.while_loop as follows:

import tensorflow as tf
tf.enable_eager_execution()

def function(array, position, times):
    array = tf.convert_to_tensor(array)
    dt = array.dtype
    ta = tf.TensorArray(dt, size=tf.size(array) - position, element_shape=())
    ta, _ = tf.while_loop(
        lambda ta, i: i < tf.size(array),
        lambda ta, i: (ta.write(i - position, tf.reduce_sum(array[i - times + 1:i + 1]) / times), i + 1),
        [ta, position])
    return ta.stack()
print(function(tf.range(30.), 10, 5).numpy())
# [ 8.  9. 10. 11. 12. 13. 14. 15. 16. 17. 18. 19. 20. 21. 22. 23. 24. 25.
#  26. 27.]

Or, even further, writing both loops instead of using tf.reduce_sum:

import tensorflow as tf
tf.enable_eager_execution()

def function(array, position, times):
    array = tf.convert_to_tensor(array)
    dt = array.dtype
    ta = tf.TensorArray(dt, size=tf.size(array) - position, element_shape=())
    ta, _ = tf.while_loop(
        lambda ta, i: i < tf.size(array),
        lambda ta, i: (ta.write(i - position, 
                                tf.while_loop(lambda s, j: j < times,
                                              lambda s, j: (s + array[i - j], j + 1),
                                              [tf.zeros((), dt), 0])[0] / times),
                       i + 1),
        [ta, position])
    return ta.stack()
print(function(tf.range(30.), 10, 5).numpy())
# [ 8.  9. 10. 11. 12. 13. 14. 15. 16. 17. 18. 19. 20. 21. 22. 23. 24. 25.
#  26. 27.]

Since you are using eager execution, you could also directly use Python loops like this:

import tensorflow as tf
tf.enable_eager_execution()

def function(array, position, times):
    array = tf.convert_to_tensor(array)
    o = []
    for loop in range(position, tf.size(array)):
        result = 0
        for k in range(times):
            result += array[loop-k]
        result /= times
        o.append(result)
    return tf.stack(o)
print(function(tf.range(30.), 10, 5).numpy())
# [ 8.  9. 10. 11. 12. 13. 14. 15. 16. 17. 18. 19. 20. 21. 22. 23. 24. 25.
#  26. 27.]

Upvotes: 1

Related Questions