Python/TensorFlow中如何相加(2,128)与(128,128)形状的张量？

问题描述 #

我有两个模型生成的损失张量，形状分别为(2, 128)和(128, 128)：

caption loss is (2, 128)
image loss is (128, 128)

损失计算代码如下：

captions_loss = keras.losses.kl_divergence(
        y_true=targets, y_pred=logits, #from_logits=True
    )

images_loss = keras.losses.kl_divergence(
        y_true=tf.transpose(targets), y_pred=tf.transpose(logits), #from_logits=True
    )

尝试直接相加时触发错误：

return (captions_loss + images_loss) / 2

解决方法 #

两个张量形状不匹配，直接相加会触发TensorFlow的形状不兼容错误，可通过以下方式处理：

1. 转为标量损失后相加（推荐） #

损失函数最终需要输出标量用于反向传播优化，因此可以对两个张量分别做全局平均或求和，得到标量后再相加：

# 对损失张量做全局平均，得到标量
captions_loss_scalar = tf.reduce_mean(captions_loss)
images_loss_scalar = tf.reduce_mean(images_loss)

# 或者用求和方式
# captions_loss_scalar = tf.reduce_sum(captions_loss)
# images_loss_scalar = tf.reduce_sum(images_loss)

return (captions_loss_scalar + images_loss_scalar) / 2

这是深度学习中损失计算的常规做法，确保最终优化的是单一标量值。

2. 通过维度扩展实现广播相加 #

如果需要保留张量维度（比如做逐元素的损失融合），可以利用TensorFlow的广播机制，先扩展张量维度让它们形状兼容：

# 将(2, 128)扩展为(2, 128, 1)
captions_loss_expanded = tf.expand_dims(captions_loss, axis=-1)
# 将(128, 128)扩展为(1, 128, 128)
images_loss_expanded = tf.expand_dims(images_loss, axis=0)

# 广播后相加，得到形状为(2, 128, 128)的张量
combined_loss = (captions_loss_expanded + images_loss_expanded) / 2

注意这种方式会生成更高维度的张量，后续需要根据业务需求决定是否进一步降维。

3. 调整损失计算时的维度匹配 #

回到损失计算逻辑，检查tf.transpose的使用是否必要。如果你的目标是让两个损失张量维度一致，可以调整转置操作或在计算KL散损时指定reduction参数，确保输出形状匹配：

# 指定reduction为NONE保留原始维度，同时调整转置逻辑
captions_loss = keras.losses.kl_divergence(
    y_true=targets, y_pred=logits,
    reduction=tf.keras.losses.Reduction.NONE
)
# 根据targets的实际形状，调整转置顺序以匹配captions_loss维度
images_loss = keras.losses.kl_divergence(
    y_true=tf.transpose(targets, perm=[1, 0]),
    y_pred=tf.transpose(logits, perm=[1, 0]),
    reduction=tf.keras.losses.Reduction.NONE
)

这种方式需要结合你的targets和logits的原始形状来调整，确保两个损失张量维度完全一致后直接相加。

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