TensorFlow中[X,10,10,1]与[X,1]张量除法的高效实现问询

我尝试将形状为[X, 10, 10, 1]的张量除以每个[10, 10, 1]子张量的和，也就是对[X, 10, 10, 1]张量与[X, 1]张量执行除法运算，但使用tf.broadcast_to时出现“形状不兼容”错误：

t1 = tf.ones([128, 10, 10, 1]) 
t1_sum = tf.reduce_sum(t1, [1, 2])  # 结果形状为[128, 1]
t1_sum_reshaped = tf.broadcast_to(t1_sum, (t1.shape))
divided = tf.math.divide(t1, t1_sum_reshaped)

当X=128时的错误提示：

Incompatible shapes: [128,1] vs. [128,10,10,1] [Op:BroadcastTo]

我目前有一个可行的临时方案，但效率极低：

t1 = tf.ones([128, 10, 10, 1])
t1_sum = tf.reduce_sum(t1, [1, 2])                                                                            # 结果形状为[128, 1]
t1_sum = tf.stack([t1_sum]*10, axis=-2)  # 结果形状为[128, 10, 1]
t1_sum = tf.stack([t1_sum]*10, axis=-2)  # 结果形状为[128, 10, 10, 1]
divided = tf.math.divide(t1, t1_sum)

请问是否存在更高效的实现方式？

高效实现方式 #

TensorFlow的广播机制无需手动复制张量，只需让两个张量维度匹配即可，以下是几种高效方案：

方法1：利用reduce_sum的keepdims参数 #

这是最简洁的方式，设置keepdims=True后，求和操作会保留被压缩的维度并设为1，直接得到可与原张量广播的形状：

t1 = tf.ones([128, 10, 10, 1])
# keepdims=True让结果形状为[128,1,1,1]
t1_sum = tf.reduce_sum(t1, [1, 2], keepdims=True)
divided = t1 / t1_sum

方法2：用tf.expand_dims扩展维度 #

手动给t1_sum添加两个单维度，使其形状变为[128,1,1,1]，触发自动广播：

t1 = tf.ones([128, 10, 10, 1])
t1_sum = tf.reduce_sum(t1, [1, 2])  # [128, 1]
# 一次性添加两个维度
t1_sum_expanded = tf.expand_dims(t1_sum, axis=[1,2])
divided = tf.math.divide(t1, t1_sum_expanded)

方法3：通过reshape调整形状 #

直接将t1_sum重塑为[128,1,1,1]，同样满足广播条件：

t1 = tf.ones([128, 10, 10, 1])
t1_sum = tf.reduce_sum(t1, [1, 2])  # [128, 1]
t1_sum_reshaped = tf.reshape(t1_sum, [-1, 1, 1, 1])
divided = tf.math.divide(t1, t1_sum_reshaped)

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