TensorFlow自定义损失函数报张量形状不匹配错误如何解决

直接触发报错的原因是循环形状配置冲突：你初始化batch_loss = 0.0时，该变量是0维标量（shape为()），但在循环内通过set_loop_options给它声明的形状不变量是一维动态张量tf.TensorShape([None])，TensorFlow做静态图编译时检测到初始形状和声明形状不匹配，直接抛出错误。
除此之外原实现还有两个会导致后续运行失败的问题：

• 用三层Python原生for循环逐像素遍历张量，在TF图模式下执行效率极低，且循环内张量形状推导极易出错
• 循环内cell_loss计算逻辑存在形状冲突：正样本分支返回的是(18,32,5)的三维张量，负样本分支返回的是0维标量，二者直接累加会触发广播错误

完全抛弃逐元素循环的写法，用TensorFlow原生向量化操作实现损失逻辑，既不会出现形状错误，运算速度也比循环实现高两个数量级，损失逻辑和原设计完全对齐：

• 网格尺寸为18*32，输出通道共5个，第0通道为置信度
• 置信度为1的正样本位置：5倍权重计算所有5个通道的L1损失
• 其余负样本位置：0.5倍权重仅计算置信度通道的L1损失
• 最终对整个batch的损失求平均

修复后的完整可运行代码如下：

backbone = tf.keras.applications.resnet50.ResNet50(include_top=False, weights=None, input_shape=INPUT_SHAPE)
x = tf.keras.layers.Conv2D(filters=5, kernel_size=3, padding='same', activation='sigmoid')(backbone.output)
model = tf.keras.Model(inputs=backbone.input, outputs=x)

def custom_loss(y_true, y_pred):
    # 提取置信度通道，形状匹配为(batch, 18, 32, 1)
    conf_true = y_true[..., 0:1]
    conf_pred = y_pred[..., 0:1]

    # 生成正负样本掩码
    pos_mask = tf.cast(tf.equal(conf_true, 1), tf.float32)
    neg_mask = 1.0 - pos_mask

    # 计算正样本损失：5倍权重，覆盖所有5个输出通道
    pos_loss = pos_mask * 5.0 * tf.abs(y_true - y_pred)
    # 计算负样本损失：0.5倍权重，仅计算置信度通道
    neg_loss = neg_mask * 0.5 * tf.abs(conf_true - conf_pred)
    # 负样本损失填充到5通道维度，和正样本损失形状对齐
    neg_loss = tf.pad(neg_loss, [[0,0], [0,0], [0,0], [0,4]], constant_values=0.0)

    # 单样本损失求和后，对batch维度求平均
    total_loss = tf.reduce_mean(tf.reduce_sum(pos_loss + neg_loss, axis=[1,2,3]))
    return total_loss

sgd = tf.keras.optimizers.SGD(momentum=0.99)
model.compile(sgd, custom_loss)

reduce_lr = tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=5)
model.fit(
    train_batch,
    validation_data = val_batch, 
    epochs = 100,
    callbacks = [reduce_lr]
)

• 移除了所有Python原生循环和错误的shape_invariants配置，从根源上避免循环内张量形状不匹配的问题
• 用掩码操作区分正负样本计算逻辑，全程保持张量维度对齐，不会出现隐式广播错误
• 用reduce_sum、reduce_mean内置算子完成损失聚合，替代手动循环累加，执行效率适配GPU/TPU并行计算

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