TensorFlow TensorArray is widely used in tf.while_loop() and tf.map_fn(). In this tutorial, we will use some examples to show you how to use it correctly.
Syntax
The initialized function of TensorArray is defined as:
__init__(
dtype,
size=None,
dynamic_size=None,
clear_after_read=None,
tensor_array_name=None,
handle=None,
flow=None,
infer_shape=True,
element_shape=None,
colocate_with_first_write_call=True,
name=None
)There are some important parameters you must notice:
size: the size of a tensorarray object
If dynamic_size = True, size can be 0
if dynamic_size = False, you will create a fixed size tensorarry, which means it can only store size elements at most.
dynamic_size: Default: False, if it is True, it means the tensorarray will change its size by elements.
clear_after_read: default: True, which means elements in this tensorarray only be allowed read one time, then be cleared. If you want to read elements in tensorarray more times, you shoud set clear_after_read = False
infer_shape: default: True, which means all elements must have the same shape.
We will use some examples to show you how to use tensorarray.
Create a fixed size tensorarray
We will use dynamic_size=False to create a fixed size tensorarray gen_o.
gen_o = tf.TensorArray(dtype=tf.float32, size = 32, dynamic_size=False, infer_shape=True)
size = 32, which means gen_o can contians 32 elements. infer_shape=True means the shape of each element in gen_o should be the same.
Create a dynamic size tensorarray
We can set dynamic_size=True to create a dynamic size tensorarray.
gen_o = tf.TensorArray(dtype=tf.float32, size = 0, dynamic_size=True, infer_shape=True)
Here, size = 0, which means the size of gen_o is 0, when you have saved an element into gen_o, the size of it will be increased.
How to write tensor into tensorarray?
We can use
write(
index,
value,
name=None
)
For example:
gen_o = gen_o.write(0, o_t)
You must notice: You only can write a tensor into the same index only once.
Meanwhile, you can not write a tensor like:
gen_o.write(0, o_t)
This way is wrong.
How to read tensor from tensorarray?
We can use
read(
index,
name=None
)
to read a tesor at the index in the tensorarray.
You must notice: if you set clear_after_read = True, the index at the index only be allowd to read once, then it will be cleaned, you can not read it again. However, if you have set clear_after_read = False, tensors in tensorarray can be read more times, but this tensor may need more memory space.
Here is an example to show how to read and write tensors
import tensorflow as tf
import numpy as np
# the shape of x1 and x2 is (3, 4)
x1 = tf.Variable(np.array([[2, 2, 3, 4], [1, 5, 3, 2] ,[6, 7, 2, 1]], dtype = np.float32), name = 'x1')
max_loop = tf.shape(x1)[0]
gen_o_1 = tf.TensorArray(dtype=tf.float32, size= max_loop + 1,
dynamic_size=False, clear_after_read = False,infer_shape=True)
init_1 = tf.Variable(np.array([0, 0, 0, 0], dtype = np.float32))
gen_o_1 = gen_o_1.write(0, init_1)
def cumulativeSum(i, gen_o_2):
next_ele = tf.nn.embedding_lookup(x1,i)
#get pre ele
pre_ele = gen_o_2.read(i)
sum = pre_ele + next_ele
i_next = i + 1
gen_o_2 = gen_o_2.write(i_next, sum)
return i_next, gen_o_2
_, gen_o_3 = tf.while_loop(cond = lambda i, _: tf.less(i, max_loop),
body = cumulativeSum,
loop_vars = (
tf.constant(0, dtype=tf.int32),
gen_o_1
)
)
v = gen_o_3.stack()
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
print(sess.run(v))
In this code, we have created a tensorarray gen_o_1 with clear_after_read = False, which means we can read tensors in it more times.
Then we have written a tensor to index = 0 first.
gen_o_1 = gen_o_1.write(0, init_1)
Meanwhile, we have read and written tensors in tf.while_loop() function.
pre_ele = gen_o_2.read(i)
sum = pre_ele + next_ele
i_next = i + 1
gen_o_2 = gen_o_2.write(i_next, sum)
Finally, run this code, you will find v is:
[[ 0. 0. 0. 0.] [ 2. 2. 3. 4.] [ 3. 7. 6. 6.] [ 9. 14. 8. 7.]]