TensorFlow中如何实现1D张量与4D张量的逐批次对应相乘计算

你可以利用TensorFlow自带的广播机制实现该逻辑，全程不需要手写循环，且是向量化运算，底层经过硬件加速，运行效率远高于手动遍历batch的实现。

实现前提 #

你提供的循环逻辑可正常运行的前提是，tensor1D[i]和tensor4D[i]满足广播条件，通常分为两种常见情况：

• 情况1：tensor1D的value维度等于4D张量的D维度，即每个通道对应一个缩放系数
• 情况2：tensor1D的value维度为1，即整个4D样本共用同一个缩放系数

两种情况都可以用以下通用代码实现：

import tensorflow as tf

tensor1D = ... # shape = [batch_size, value]
tensor4D = ... # shape = [batch_size, L, W, D]

# 给1D张量扩展两个维度，适配4D张量的形状，得到shape为[batch_size, 1, 1, value]的张量
expanded_1d = tf.expand_dims(tf.expand_dims(tensor1D, axis=1), axis=1)
# 直接相乘，TensorFlow会自动完成维度广播
result = expanded_1d * tensor4D

等价实现 #

你也可以用tf.reshape实现维度扩展，效果完全一致：

expanded_1d = tf.reshape(tensor1D, (-1, 1, 1, tensor1D.shape[-1]))
result = expanded_1d * tensor4D

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