Python tf.TensorArray用法及代码示例

类包装 dynamic-sized、per-time-step、write-once 张量数组。

用法

tf.TensorArray(
    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
)

此类旨在与动态迭代原语一起使用，例如 while_loop 和 map_fn 。它通过特殊的"flow"控制流依赖支持渐变back-propagation。

示例1：简单的阅读和写作。

ta = tf.TensorArray(tf.float32, size=0, dynamic_size=True, clear_after_read=False)
ta = ta.write(0, 10)
ta = ta.write(1, 20)
ta = ta.write(2, 30)

ta.read(0)
<tf.Tensor:shape=(), dtype=float32, numpy=10.0>
ta.read(1)
<tf.Tensor:shape=(), dtype=float32, numpy=20.0>
ta.read(2)
<tf.Tensor:shape=(), dtype=float32, numpy=30.0>
ta.stack()
<tf.Tensor:shape=(3,), dtype=float32, numpy=array([10., 20., 30.],
dtype=float32)>

示例 2：循环写入然后返回的斐波那契序列算法。

@tf.function
def fibonacci(n):
  ta = tf.TensorArray(tf.float32, size=0, dynamic_size=True)
  ta = ta.unstack([0., 1.])

  for i in range(2, n):
    ta = ta.write(i, ta.read(i - 1) + ta.read(i - 2))

  return ta.stack()

fibonacci(7)
<tf.Tensor:shape=(7,), dtype=float32,
numpy=array([0., 1., 1., 2., 3., 5., 8.], dtype=float32)>

示例 3：与 tf.Variable 交互的简单循环。

v = tf.Variable(1)
@tf.function
def f(x):
  ta = tf.TensorArray(tf.int32, size=0, dynamic_size=True)
  for i in tf.range(x):
    v.assign_add(i)
    ta = ta.write(i, v)
  return ta.stack()
f(5)
<tf.Tensor:shape=(5,), dtype=int32, numpy=array([ 1,  2,  4,  7, 11],
dtype=int32)>