TensorFlow中按每3行求和生成新矩阵的技术咨询

没问题，这个需求用TensorFlow的内置操作就能轻松实现，而且完全支持可微分计算，完美适配你的计算图需求。我给你一步步拆解：

核心思路 #

我们可以通过重塑矩阵形状把原矩阵的行按每3行分组，然后在组内对行维度求和，这样就能得到每3行的累加结果。所有用到的操作都是TensorFlow原生支持的，天然可微分。

代码实现 #

import tensorflow as tf

# 你的示例输入
input_matrix = tf.constant([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12]])

# 获取矩阵的行数和列数（用tf.shape保证动态计算，适配任意大小的矩阵）
num_rows = tf.shape(input_matrix)[0]
num_cols = tf.shape(input_matrix)[1]

# 重塑矩阵：将原矩阵转为 [组数, 每组3行, 列数] 的形状
# 这里的-1会让TensorFlow自动计算组数（总行数/3）
reshaped_matrix = tf.reshape(input_matrix, (-1, 3, num_cols))

# 沿着每组的行维度（axis=1）求和，得到最终结果
output_matrix = tf.reduce_sum(reshaped_matrix, axis=1)

# 验证结果
print(output_matrix.numpy())
# 输出正好符合预期：
# [[ 9 12]
#  [27 30]]

关键操作解释 #

• tf.reshape：把原矩阵的行重新分组，比如6行2列的矩阵会被转为2组，每组3行2列。这个操作只是改变数据的视图，不会修改原始数据，计算效率很高。
• tf.reduce_sum：对每组内的3行进行逐列求和，axis=1指定了求和的维度是每组的行维度，最终得到每组一行的结果矩阵。

处理非3的倍数行数的情况（可选） #

如果你的原始矩阵行数不是3的整数倍，比如7行，可以先通过tf.pad补零到最近的3的倍数，再执行上述操作：

# 假设输入是7行2列
input_matrix = tf.constant([[1,2],[3,4],[5,6],[7,8],[9,10],[11,12],[13,14]])
num_rows = tf.shape(input_matrix)[0]
# 计算需要补的行数
pad_rows = tf.math.mod(-num_rows, 3)
# 补零操作
padded_matrix = tf.pad(input_matrix, [[0, pad_rows], [0, 0]])
# 后续步骤和之前一样
reshaped = tf.reshape(padded_matrix, (-1,3,tf.shape(input_matrix)[1]))
output = tf.reduce_sum(reshaped, axis=1)
print(output.numpy())
# 输出：
# [[ 9 12]
#  [27 30]
#  [13 14]]

所有这些操作都是完全可微分的，不会破坏你的TensorFlow计算图反向传播流程。

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