TensorFlow 2.13下Mask R-CNN自定义损失函数引发KerasTensor TypeError问题求助

大家好，我在使用TensorFlow 2.13版本运行Mask R-CNN模型时遇到了麻烦：模型在训练前的环节都正常，但启动训练时就触发了KerasTensor相关的TypeError，我推测问题出在自定义损失函数的写法上。

以下是我的代码细节，麻烦大家帮忙排查一下：

模型编译部分 #

def compile(self, learning_rate, momentum):
    """Compiles the Mask R-CNN model by adding losses and metrics."""
    # Check if the model is built before compiling
    if not hasattr(self, 'keras_model') or self.keras_model is None:
        print("⚠️ Model not built. Building now...")
        self.keras_model = self.build(mode="training", config=self.config)

    # Optimizer setup
    optimizer = tf.keras.optimizers.SGD(
        learning_rate=learning_rate, 
        momentum=momentum,
        clipnorm=self.config.GRADIENT_CLIP_NORM
    )
    
    loss_names = [
        "rpn_class_loss", "rpn_bbox_loss",
        "mrcnn_class_loss", "mrcnn_bbox_loss", "mrcnn_mask_loss"
    ]

    print(f"✅ Available layers: {[layer.name for layer in self.keras_model.layers]}")
    
    for loss_name in loss_names:
        try:
            loss_layer = self.keras_model.get_layer(loss_name)
            loss_value = tf.reduce_mean(loss_layer.output)
            print(f"✅ Loss layer {loss_name} found in model: {loss_value}")
            #loss_value = tf.convert_to_tensor(loss_value, dtype=tf.float32)
            @tf.function
            def compute_loss():
                return tf.reduce_mean(loss_layer.output)

            loss_value = compute_loss()
            print(f"✅ Loss layer {loss_name} converted from KerasTensor to normal tensor: {loss_value}")
            
            self.keras_model.add_loss(lambda: loss_value)
            print(f"✅ Added loss: {loss_name}")
        except ValueError:
            print(f"❌ Warning: Loss layer {loss_name} not found in model!")
    
    config_dict = {k: v for k, v in self.config.__dict__.items() if not k.startswith('__')}       
    # Compile the model
    self.keras_model.compile(
        optimizer=optimizer,
        loss=None,   # Mask R-CNN defines its own loss layers
        run_eagerly=True,
        experimental_run_tf_function=False
    )

    print("✅ Model compiled successfully!")

模型构建中添加损失的代码片段 #

rpn_class_loss = RPNClassLoss(name="rpn_class_loss")([input_rpn_match, rpn_class_logits])
self.add_loss(tf.reduce_mean(rpn_class_loss))
rpn_bbox_loss = RPNBBoxLoss(name="rpn_bbox_loss")([config, input_rpn_bbox, input_rpn_match, rpn_bbox])
self.add_loss(tf.reduce_mean(rpn_bbox_loss))
mrcnn_class_loss = MRCNNClassLoss(name="mrcnn_class_loss")([target_class_ids, mrcnn_class_logits, active_class_ids])
self.add_loss(tf.reduce_mean(mrcnn_class_loss))
mrcnn_bbox_loss = MRCNNBBoxLoss(name="mrcnn_bbox_loss")([target_bbox, target_class_ids, mrcnn_bbox])
self.add_loss(tf.reduce_mean(mrcnn_bbox_loss))
mrcnn_mask_loss = MRCNNMaskLoss(name="mrcnn_mask_loss")([target_mask, target_class_ids, mrcnn_mask])
self.add_loss(tf.reduce_mean(mrcnn_mask_loss))

自定义损失函数类实现 #

SmoothL1Loss #

class SmoothL1Loss(Layer):
    @tf.function
    def call(self, inputs):
        """Computes Smooth-L1 loss."""
        y_true, y_pred = inputs  # Expecting two tensors: [batch, N, 4]

        diff = tf.abs(y_true - y_pred)
        less_than_one = tf.cast(tf.less(diff, 1.0), tf.float32)

        loss = (less_than_one * 0.5 * tf.square(diff)) + ((1.0 - less_than_one) * (diff - 0.5))
        return tf.reduce_mean(loss)  # ✅ Ensures numerical value

RPNClassLoss #

class RPNClassLoss(Layer):
    @tf.function
    def call(self, inputs):
        rpn_match, rpn_class_logits = inputs  # Shapes: [batch, anchors, 1] and [batch, proposals, 2, 2]

        rpn_match = tf.squeeze(rpn_match, -1)  # Shape: [batch, anchors]

        # Convert -1/+1 labels to 0/1 format
        anchor_class = tf.where(rpn_match == 1, 1, 0)  # Shape: [batch, anchors]

        # ✅ Ensure indices do not exceed the available range
        max_proposals = tf.shape(rpn_class_logits)[1]  # Get the actual number of proposals
        valid_indices = tf.where(
            tf.logical_and(tf.not_equal(rpn_match, 0), tf.range(tf.shape(rpn_match)[1]) < max_proposals)
        )  # Shape: [num_valid, 2]

        # ✅ Gather only valid anchors
        rpn_class_logits = tf.gather_nd(rpn_class_logits, valid_indices)  # Shape: [num_valid, 2, 2]
        anchor_class = tf.gather_nd(anchor_class, valid_indices)  # Shape: [num_valid]

        #  **Fix: Ensure correct shape for labels**
        rpn_class_logits = tf.reshape(rpn_class_logits, [-1, tf.shape(rpn_class_logits)[-1]])  # [num_valid, num_classes]
        anchor_class = tf.reshape(anchor_class, [-1])  # [num_valid]

        # ✅ Ensure labels and logits match in shape
        min_size = tf.minimum(tf.shape(anchor_class)[0], tf.shape(rpn_class_logits)[0])
        anchor_class = anchor_class[:min_size]
        rpn_class_logits = rpn_class_logits[:min_size]

        # ✅ Compute loss
        loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=anchor_class, logits=rpn_class_logits)

        return tf.reduce_mean(loss)  # Ensure numerical loss

RPNBBoxLoss #

class RPNBBoxLoss(Layer):
    @tf.function
    def call(self, inputs):
        config, target_bbox, rpn_match, rpn_bbox = inputs  # Expecting 4 inputs

        rpn_match = tf.squeeze(rpn_match, -1)  # Shape: [batch, anchors]
        indices = tf.where(tf.equal(rpn_match, 1))  # Select positive anchors

        # ✅ Convert indices[:, 1] to int32 to match max_indices
        indices = tf.cast(indices, tf.int32)

        max_indices = tf.shape(target_bbox)[1]  # Keep this as int32
        valid_mask = tf.less(indices[:, 1], max_indices)  # Ensure valid range
        valid_indices = tf.boolean_mask(indices, valid_mask)  # Remove invalid indices

        # Gather only the positive anchors' bbox values
        rpn_bbox = tf.gather_nd(rpn_bbox, valid_indices)  # Shape: [num_pos, 2, 4]
        target_bbox = tf.gather_nd(target_bbox, valid_indices)  # Shape: [num_pos, 4]
        
        # ✅ Fix shape mismatch by selecting first prediction (or averaging)
        rpn_bbox = tf.reduce_mean(rpn_bbox, axis=1)  # Shape: [num_pos, 4]

        # Compute loss safely using tf.cond to avoid NaNs
        num_samples = tf.shape(target_bbox)[0]
        loss = tf.cond(
            num_samples > 0,
            lambda: tf.reduce_mean(SmoothL1Loss()([target_bbox, rpn_bbox])),
            lambda: tf.constant(0.0, dtype=tf.float32)
        )
        
        return loss  # Ensure a valid numerical output

MRCNNClassLoss #

class MRCNNClassLoss(Layer):
    @tf.function
    def call(self, inputs):
        target_class_ids, pred_class_logits, active_class_ids = inputs

        loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=target_class_ids, logits=pred_class_logits)
        return tf.reduce_mean(loss)

MRCNNBBoxLoss #

class MRCNNBBoxLoss(Layer):
    @tf.function
    def call(self, inputs):
        target_bbox, target_class_ids, pred_bbox = inputs

        # Reshape for gathering
        target_class_ids = tf.reshape(target_class_ids, [-1])  # Shape: [batch * num_rois]
        target_bbox = tf.reshape(target_bbox, [-1, 4])  # Shape: [batch * num_rois, 4]
        pred_bbox = tf.reshape(pred_bbox, [-1, tf.shape(pred_bbox)[2], 4])  # Shape: [batch * num_rois, num_classes, 4]

        # Get positive ROIs
        positive_roi_ix = tf.where(target_class_ids > 0)

        # ✅ Ensure positive_roi_ix is not empty
        has_valid_rois = tf.greater(tf.shape(positive_roi_ix)[0], 0)

        def compute_loss():
            """Inner function for computing loss, wrapped for TensorFlow execution."""
            positive_roi_ix_flat = tf.reshape(positive_roi_ix, [-1])  # Flatten indices
            positive_class_ids = tf.gather(target_class_ids, positive_roi_ix_flat)
            
            positive_roi_ix_flat = tf.cast(positive_roi_ix_flat, tf.int32)
            positive_class_ids = tf.cast(positive_class_ids, tf.int32)

            indices = tf.stack([positive_roi_ix_flat, positive_class_ids], axis=1)
            pred_bbox_selected = tf.gather_nd(pred_bbox, indices)

            loss = tf.reduce_mean(SmoothL1Loss()([target_bbox, pred_bbox_selected]))  
            return tf.cast(loss, tf.float32)  # ✅ Ensure TensorFlow tensor output

        # Use `tf.cond` to handle empty cases
        loss = tf.cond(
            has_valid_rois,
            compute_loss,  # Call compute_loss if has_valid_rois is True
            lambda: tf.constant(0.0, dtype=tf.float32)  # Return 0.0 if has_valid_rois is False
        )

        return loss

MRCNNMaskLoss #

class MRCNNMaskLoss(Layer):
    @tf.function
    def call(self, inputs):
        target_masks, target_class_ids, pred_masks = inputs

        # Flatten target_class_ids to be 1D: (batch * num_rois,)
        target_class_ids = tf.reshape(target_class_ids, [-1])
        
        # Flatten target_masks
        mask_shape = tf.shape(target_masks)
        target_masks = tf.reshape(target_masks, [-1, mask_shape[2], mask_shape[3]])
        
        # Reshape pred_masks
        pred_shape = tf.shape(pred_masks)
        pred_masks = tf.reshape(pred_masks, [-1, pred_shape[2], pred_shape[3], pred_shape[4]])
        
        # Permute predicted masks to [N, num_classes, height, width]
        pred_masks = tf.transpose(pred_masks, [0, 3, 1, 2])
        
        # Get indices of positive ROIs
        positive_ix = tf.where(target_class_ids > 0)[:, 0]
        
        # Ensure dtype consistency
        positive_ix = tf.cast(positive_ix, tf.int32)
        positive_class_ids = tf.cast(tf.gather(target_class_ids, positive_ix), tf.int32)
        
        # Stack indices for pred_masks gathering
        indices = tf.stack([positive_ix, positive_class_ids], axis=1)  # ✅ Fix dtype issue
        
        # Gather the masks (predicted and true) that contribute to loss
        y_true = tf.gather(target_masks, positive_ix)
        y_pred = tf.gather_nd(pred_masks, indices)
        
        # Compute binary cross entropy loss
        num_samples = tf.shape(y_true)[0]
        loss = tf.cond(
            tf.greater(num_samples, 0),  # Check if num_samples > 0
            lambda: tf.reduce_mean(tf.keras.losses.binary_crossentropy(y_true, y_pred)),
            lambda: tf.constant(0.0, dtype=tf.float32)
        )
        return loss

备注：内容来源于stack exchange，提问作者Naimur Hassan