使用Keras和Tensorflow搭建自定义IoU层出现维度不兼容错误如何解决

错误原因 #

• 首先是tf.map_fn的使用逻辑错误：当你给tf.map_fn传入多组张量作为迭代对象时，它要求所有张量的第0维长度必须完全一致，会按位取每一组元素传入映射函数。你现在传入的gt_boxes第0维长度是3，anchors第0维长度是10，维度不匹配就触发了这个报错。
• 其次是你用多层嵌套循环+Python原生逻辑写IoU计算，在TF1.x关闭eager执行的场景下，这类原生Python逻辑不会被转换为静态计算图节点，就算解决了维度问题后续也会跑不通。
• 最后你嵌套调用tf.map_fn的时候参数顺序传错了，multiple_IoU里的tf.map_fn参数顺序不符合API要求。

修复方案 #

你要算两组框的两两IoU，完全不需要用循环/map_fn，直接用TF的广播机制就行，速度快还不会有维度问题，修复后的代码如下：

import sys
import numpy as np 
from tensorflow.keras.layers import Layer
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()

class IoULayer(Layer):
    def __init__(self, **kwargs):
        super(IoULayer, self).__init__(**kwargs)
        
    def call(self, inputs):
        '''
        inputs[0]: ground truth boxes  维度 [N, 4]
        inputs[1]: anchors boxes       维度 [M, 4]
        输出: 两两IoU 维度 [N, M]
        '''    
        gt_boxes = inputs[0]
        anchors = inputs[1]
        
        # 扩展维度做广播
        gt_expand = tf.expand_dims(gt_boxes, axis=1)  # [N, 1, 4]
        anchor_expand = tf.expand_dims(anchors, axis=0)  # [1, M, 4]
        
        # 计算交集坐标
        x_left = tf.maximum(gt_expand[..., 0], anchor_expand[..., 0])
        y_top = tf.maximum(gt_expand[..., 1], anchor_expand[..., 1])
        x_right = tf.minimum(gt_expand[..., 2], anchor_expand[..., 2])
        y_bottom = tf.minimum(gt_expand[..., 3], anchor_expand[..., 3])
        
        # 计算交集面积
        intersect_w = tf.maximum(x_right - x_left, 0.0)
        intersect_h = tf.maximum(y_bottom - y_top, 0.0)
        intersect_area = intersect_w * intersect_h
        
        # 计算各自面积
        gt_area = (gt_expand[..., 2] - gt_expand[..., 0]) * (gt_expand[..., 3] - gt_expand[..., 1])
        anchor_area = (anchor_expand[..., 2] - anchor_expand[..., 0]) * (anchor_expand[..., 3] - anchor_expand[..., 1])
        
        # 计算IoU，避免除零
        union_area = gt_area + anchor_area - intersect_area
        iou = tf.where(union_area > 0, intersect_area / union_area, tf.zeros_like(intersect_area))
        
        return iou

额外说明 #

如果你确实需要保留map_fn的写法（不推荐，性能差很多），需要把call里的map_fn去掉，直接调用计算逻辑，因为你是要算两个张量所有元素的两两匹配，不是按位迭代两个张量的元素。

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