Keras后端switch结合tf.where无法按预期工作的问题

问题根源 #

你原本用tf.where提取one-hot张量中值为1的坐标，再取s2维度索引，但当预测全为0时，tf.where返回空张量，导致K.switch因分支形状不匹配报错，且这种稀疏提取的方式无法直接生成你需要的[batch_size,S1,S2,1]稠密张量。

正确实现思路 #

改用稠密张量操作直接生成目标形状的结果，彻底避免依赖稀疏的tf.where：

1. 提取目标one-hot子张量
2. 判断每个空间位置是否全为0
3. 提取每个位置的one-hot通道索引（对应你需要的特定范围值）
4. 对全0位置替换为0，最终输出符合形状要求的张量

代码实现 #

# 提取one-hot部分的张量，形状[batch_size, S1, S2, 12]
one_hot_box = y_pred[..., C+1:C+13]

# 判断每个空间位置是否全为0，形状[batch_size, S1, S2, 1]
is_all_zero = tf.reduce_all(tf.equal(one_hot_box, 0), axis=-1, keepdims=True)

# 提取每个位置的one-hot通道索引，形状[batch_size, S1, S2, 1]
# 注：如果one-hot是严格单热的，这里就是对应1所在的通道
channel_indices = tf.argmax(one_hot_box, axis=-1, keepdims=True)

# 替换全0位置的索引为0，最终输出形状[batch_size, S1, S2, 1]
where_box1 = tf.where(is_all_zero, tf.zeros_like(channel_indices), channel_indices)

额外说明 #

如果你原本的tf.where(...)[...,2]是想获取空间位置的s2索引而非通道索引，那逻辑本身就不符合生成稠密张量的需求，此时需要重新梳理IOU计算的输入要求，确保操作对象是稠密张量而非稀疏坐标集合。

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