在 for 循环中存储张量数组答案

【问题标题】：Store array of tensors in for loop在 for 循环中存储张量数组
【发布时间】：2022-08-16 04:10:19
【问题描述】：

我有一个 @tf.function 装饰函数。在函数内部，我想从分布中提取并多次计算某个值（让我们将此过程称为 f(x)）。

如何在 TensorFlow 2.0 中做到这一点？我不能使用 numpy 数组，因为我想使用 @tf.function 装饰器。

一个 numpy 实现看起来像：

reps = 4
store = np.zeros((n, reps))
for i in range(reps):
  store[:, i] = f(x) #f(x) is shape (n,)

然后目标是计算store 的行均值。

这应该很容易，但我无法弄清楚如何去做！

标签： python tensorflow tensorflow2.0 immutability tf.keras

【解决方案1】：

可能是这样的：

import tensorflow as tf

def f():
  return tf.random.normal((10,))

@tf.function
def store_this():
  reps = 4
  n = 10
  store = tf.zeros((n, reps))
  values = [f() for _ in range(reps)]
  indices = tf.stack([tf.tile(tf.range(n), multiples=[reps]), tf.repeat(tf.range(reps), repeats=n)], axis=-1)
  return tf.tensor_scatter_nd_update(store, indices, tf.reshape(values, [-1]))

store_this()

【讨论】：

谢谢，我会用tf.tensor_scatter_nd_update。与 pytorch 相比，在 tensorflow 中执行此操作似乎要麻烦得多，但它确实有效。

【解决方案2】：

如果 f 将一维张量作为输入，这是一个较短的替代方案：

@tf.function
def f(x):
  return tf.random.normal((10,))
x = tf.constant([1.0, 2.0])
reps = 4

def store_this(fp, x, reps):
    return tf.transpose(tf.map_fn(fp, tf.tile(tf.expand_dims(x, 0),[reps,1])))
store_this(f, x, reps)

【讨论】：