如何对大型矩阵的不同区间进行批量reduce_sum运算？

我来帮你搞定这个TensorFlow批量区间求和的问题，不用tf.py_func也能在GPU上顺畅运行～

核心思路 #

我们可以通过生成区间掩码的方式，把每个样本需要求和的行标记出来，再用矩阵乘法实现批量加权求和——全程都是原生TensorFlow操作，完美兼容GPU，不会有tf.py_func的性能或兼容性问题。

完整实现代码 #

import tensorflow as tf
tf.enable_eager_execution()

emb = tf.ones([100,16])
start_pos = tf.constant([1,2])
end_pos = tf.constant([11,31])

# 获取各维度的形状信息
batch_size = tf.shape(start_pos)[0]
seq_len = tf.shape(emb)[0]
hidden_size = tf.shape(emb)[1]

# 生成形状为(1, seq_len)的全局位置序列
pos_range = tf.range(seq_len, dtype=tf.int32)[tf.newaxis, :]
# 将起始/结束位置扩展为(batch_size, 1)，方便广播比较
start_expanded = start_pos[:, tf.newaxis]
end_expanded = end_pos[:, tf.newaxis]

# 生成掩码：每个位置是否属于当前样本的求和区间（包含两端）
mask = tf.logical_and(pos_range >= start_expanded, pos_range <= end_expanded)
# 转换为浮点型，用于后续加权求和
mask_float = tf.cast(mask, dtype=emb.dtype)

# 通过矩阵乘法实现批量求和：mask_float (2,100) × emb (100,16) → (2,16)
result = tf.matmul(mask_float, emb)

# 验证结果
print(f"输出形状：{result.shape}")
print(f"求和结果：\n{result.numpy()}")

代码说明 #

1. 掩码生成：利用TensorFlow的广播机制，把每个样本的起始/结束位置和全局位置序列做比较，快速生成每个样本对应的有效行掩码。
2. 批量求和：用矩阵乘法替代循环求和，既简洁又能利用GPU的并行计算能力，效率拉满。
3. 灵活性：如果你的区间是左闭右开的需求，只需要把pos_range <= end_expanded改成pos_range < end_expanded即可。

这个方法完全规避了tf.py_func的问题，不管是CPU还是GPU环境都能稳定运行，而且计算效率很高，适合大规模数据场景～

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