如何在TensorFlow中高效实现滑动窗口式分组？

需求回顾 #

给定n维张量，指定某一维度为时间维度，需将该维度替换为[n_groups, window_size]两个新维度：

• window_size为指定的滑动窗口大小
• n_groups是时间维度上所有连续窗口的数量，正确计算应为n_periods - window_size + 1（注：原描述中n_groups = n_periods - window_size为笔误，对应示例代码的循环逻辑应为range(len(inputs) - window_size + 1)）

传统方法的问题 #

用Python循环切片再tf.stack的方式（如下），在时间维度较长时会生成大量图操作，导致计算图臃肿、效率低下：

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

TensorFlow内置高效解决方案：tf.signal.frame #

TensorFlow提供了**tf.signal.frame**函数，专门用于高效生成滑动窗口（分帧）操作，底层实现避免了循环生成大量操作的问题，完全适配你的需求。

核心参数说明

• frame_length: 滑动窗口的大小（即你的window_size）
• frame_step: 窗口滑动的步长，这里要取所有连续窗口，所以设为1
• axis: 指定需要拆分的时间维度的索引

代码示例

假设输入张量形状为[batch_size, time_steps, features]，时间维度为axis=1：

import tensorflow as tf

# 构造示例张量：3个样本，10个时间步，5个特征
inputs = tf.random.normal((3, 10, 5))
window_size = 3
time_axis = 1

# 生成滑动窗口
output = tf.signal.frame(
    signal=inputs,
    frame_length=window_size,
    frame_step=1,
    axis=time_axis
)

print(f"原张量形状: {inputs.shape}")
print(f"输出张量形状: {output.shape}")

输出结果：

原张量形状: (3, 10, 5)
输出张量形状: (3, 8, 3, 5)

可以看到，原时间维度（长度10）被拆分为n_groups=8和window_size=3两个维度，完全符合需求。

适配不同维度位置

如果时间维度是其他位置（比如第一个维度或最后一个维度），只需调整axis参数即可：

# 时间维度为第0位的情况，输入形状(10, 5)
inputs = tf.random.normal((10, 5))
output = tf.signal.frame(inputs, frame_length=3, frame_step=1, axis=0)
print(output.shape)  # (8, 3, 5)

内容的提问来源于stack exchange，提问作者Alex Pilafian