如何在TensorFlow中用向量化操作实现三维张量的二维滑动窗口？

你可以直接用TensorFlow内置的tf.image.extract_patches函数来实现，完全不需要循环或依赖NumPy，这是专门为滑动窗口提取设计的向量化操作，完美匹配你的需求。

具体实现步骤 #

假设你的三维张量形状是[height, width, channels]（每个位置的n元素数组对应channels=n），以提取3×3滑动窗口为例：

1. 适配输入维度
   tf.image.extract_patches默认处理4D张量（格式为[batch, height, width, channels]），所以先给你的三维张量加一个batch维度：

# 示例：创建一个20×20、每个位置是5元素数组的三维张量
input_tensor = tf.random.normal(shape=[20, 20, 5])
input_4d = tf.expand_dims(input_tensor, axis=0)

2. 配置滑动窗口参数

• sizes：窗口尺寸，对应4D张量的四个维度，设为[1, 3, 3, 1]表示batch维度取1个样本，空间维度取3×3窗口，通道维度完整保留
• strides：窗口滑动步长，[1, 1, 1, 1]表示逐像素滑动，若需要间隔滑动，调整中间两个数值即可
• padding：填充方式，"VALID"表示窗口必须完全落在张量内部（输出尺寸会缩小），"SAME"表示自动填充使输出尺寸和原张量空间维度一致

3. 提取滑动窗口

patches = tf.image.extract_patches(
    images=input_4d,
    sizes=[1, 3, 3, 1],
    strides=[1, 1, 1, 1],
    rates=[1, 1, 1, 1],
    padding="VALID"
)

4. 调整输出形状
   提取后的patches形状是[1, output_h, output_w, 3*3*channels]，可以把最后一维reshape成直观的窗口结构：

# 转换为[output_h, output_w, 3, 3, channels]，方便查看每个3×3窗口的细节
patches_reshaped = tf.reshape(
    patches,
    [-1, tf.shape(patches)[1], tf.shape(patches)[2], 3, 3, input_tensor.shape[-1]]
)
# 去掉多余的batch维度
patches_reshaped = tf.squeeze(patches_reshaped, axis=0)

灵活调整窗口大小 #

如果需要提取5×5这类其他尺寸的窗口，只需修改两处：

• 将extract_patches的sizes参数改为[1,5,5,1]
• reshape时把最后一维的3,3换成5,5即可

验证示例 #

import tensorflow as tf

# 构造测试张量：10×10，每个位置是3元素数组
input_tensor = tf.range(10*10*3, dtype=tf.float32)
input_tensor = tf.reshape(input_tensor, [10, 10, 3])

# 扩展为4D张量
input_4d = tf.expand_dims(input_tensor, axis=0)

# 提取3×3滑动窗口
patches = tf.image.extract_patches(
    images=input_4d,
    sizes=[1,3,3,1],
    strides=[1,1,1,1],
    rates=[1,1,1,1],
    padding="VALID"
)

# 调整为直观的窗口形状
patches_reshaped = tf.reshape(patches, [8, 8, 3, 3, 3])

# 验证提取结果和原张量对应区域一致
print("原张量左上角3×3区域：")
print(input_tensor[:3, :3, :])
print("\n提取的第一个滑动窗口：")
print(patches_reshaped[0, 0, :, :, :])

运行后会看到两者完全匹配，说明提取正确。

内容的提问来源于stack exchange，提问作者Luiz Doleron