TensorFlow：对3D张量的子张量执行reduce_mean并拼接结果

嘿，这个问题我熟！咱们来一步步搞定它——你需要把一个形状为(X, Y, Z)的3D张量，沿着X轴分成N个等长的块（因为N能整除X），每个块求均值后再拼接成(N, Y, Z)的张量对吧？

最简洁高效的实现方法：用reshape+reduce_mean #

这种方法利用TensorFlow的形状重塑，避免了拆分和拼接的额外开销，是最优解：

import tensorflow as tf

def compute_block_mean(input_tensor, N):
    # 动态获取运行时的张量形状（因为X、Y是运行时确定的）
    X, Y, Z = tf.shape(input_tensor)[0], tf.shape(input_tensor)[1], tf.shape(input_tensor)[2]
    # 把原张量重塑为 (N, X//N, Y, Z)：把X轴拆成N个长度为X/N的块
    reshaped_tensor = tf.reshape(input_tensor, (N, X // N, Y, Z))
    # 对每个块的内部维度（也就是第二个轴，X//N的维度）求均值
    block_means = tf.reduce_mean(reshaped_tensor, axis=1)
    return block_means

举个例子验证 #

比如当N=2，输入张量形状是(4, 3, 2)（X=4能被2整除）：

1. 重塑后变成(2, 2, 3, 2)，相当于把原X轴的4个元素分成了2个各含2个元素的块
2. 对第二个轴（索引1）求均值后，就得到每个块的均值，最终形状是(2, 3, 2)，完全符合你的需求！

另一种更直观的实现：split+stack #

如果你觉得reshape的逻辑有点绕，也可以用拆分+逐个求均值+堆叠的方式，逻辑更直白：

def compute_block_mean_split(input_tensor, N):
    # 沿着X轴把张量拆分成N个等长的子张量，每个子张量形状是(X//N, Y, Z)
    tensor_splits = tf.split(input_tensor, num_or_size_splits=N, axis=0)
    # 对每个子张量在X轴上求均值，得到形状为(Y, Z)的张量
    mean_tensors = [tf.reduce_mean(split, axis=0) for split in tensor_splits]
    # 把所有均值张量堆叠起来，得到(N, Y, Z)的结果
    block_means = tf.stack(mean_tensors, axis=0)
    return block_means

注意事项 #

• 一定要用tf.shape(input_tensor)来获取动态形状，因为X、Y是运行时才确定的，而tensor.shape是静态形状，可能无法正确获取运行时的实际维度。
• 两种方法都能得到相同的结果，但reshape的方法性能更好，因为它不需要生成多个中间张量，直接通过形状变换完成操作。

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