在 TensorFlow 中创建“滚动窗口”类型分组的有效方法

Question

假设您有一个 n 维张量，其中一个维度对应于时间。

我想做的是：给定一些 integer window_size ，我想用两个新维度[..., n_groups, window_size]替换我的时间维度。 其中n_groups代表时间维度上大小为window_size的所有可能分组。 因此，如果我们从大小为n_periods的时间维度开始，那么n_groups最终应该是n_periods - window_size 。

使用传统的“pythonic”循环和切片可以很容易地完成所有这些，例如：

stacked = tf.stack([inputs[i:i+window_size] for i in range(len(inputs) - window_size + 1)], axis=0)

但是，如果时间维度很长，则会产生数量惊人的图形操作。 我想知道是否没有内置的 TensorFlow function 可以帮助我更有效地完成这个相对简单的任务......

“滚动窗口分组”的想法如此普遍，以至于 Pandas 项目有一个非常复杂且规模庞大的 API来处理这种特殊情况。 我本以为 TensorFlow 也会包含这样的实用程序。

Answer 1

考虑关于 map_fn 的 tf 文档：

“map_fn 会将 fn 使用的操作应用于 elems 的每个元素，从而导致 O(elems.shape[0]) 总操作数。由于 map_fn 可以并行处理元素这一事实在一定程度上缓解了这一点。但是，使用map_fn 通常仍然比使用矢量化操作表示的等效转换效率低。”

给定输入张量，您可以尝试以下方法：

input_tensor = tf.range([10])

# <tf.Tensor: shape=(10,), dtype=int32, numpy=array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=int32)>

转换为方阵：

 res = tf.repeat(tf.expand_dims(input_tensor, 0), input_tensor.shape[0], axis = 0)

  # array([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
  # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
  # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
  # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
  # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
  # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
  # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
  # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
  # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
  # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]], dtype=int32)>

然后在这个张量上应用 map_fn，在输入中包含一个具有负值的范围向量：

elements = tf.range(10, dtype=tf.int32) * -1
w,_ = tf.map_fn(lambda x: (tf.roll(x[0], x[1], axis=0), x[1]), (res, elements), dtype=(tf.int32, tf.int32))

这会将元素排（左）为：

#<tf.Tensor: shape=(10, 10), dtype=int32, numpy=
#array([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
#       [1, 2, 3, 4, 5, 6, 7, 8, 9, 0],
#       [2, 3, 4, 5, 6, 7, 8, 9, 0, 1],
#       [3, 4, 5, 6, 7, 8, 9, 0, 1, 2],
#       [4, 5, 6, 7, 8, 9, 0, 1, 2, 3],
#       [5, 6, 7, 8, 9, 0, 1, 2, 3, 4],
#       [6, 7, 8, 9, 0, 1, 2, 3, 4, 5],
#       [7, 8, 9, 0, 1, 2, 3, 4, 5, 6],
#       [8, 9, 0, 1, 2, 3, 4, 5, 6, 7],
#       [9, 0, 1, 2, 3, 4, 5, 6, 7, 8]], dtype=int32)>

最后，使用张量切片获取尽可能多的元素，例如：

 window = 8
 tf.slice(w, [0, 0], [(w.shape[0] - window) + 1, window])

给出：

#<tf.Tensor: shape=(3, 8), dtype=int32, numpy=
#array([[0, 1, 2, 3, 4, 5, 6, 7],
#       [1, 2, 3, 4, 5, 6, 7, 8],
#       [2, 3, 4, 5, 6, 7, 8, 9]], dtype=int32)>

对于 window = 4

window = 4
tf.slice(w, [0, 0], [(w.shape[0] - window) + 1, window])

给出：

#array([[0, 1, 2, 3],
#   [1, 2, 3, 4],
#   [2, 3, 4, 5],
#   [3, 4, 5, 6],
#   [4, 5, 6, 7],
#   [5, 6, 7, 8],
#   [6, 7, 8, 9]], dtype=int32)>

尝试。 将其转换成 tf 图，看看它是否比普通的 python 循环具有更好的性能。