自定义Faster R-CNN的RPN损失函数遇无梯度问题求助

梯度无传播的核心原因 #

你的代码大量使用Python原生操作脱离TensorFlow计算图追踪机制，直接切断梯度流：

• 将Tensor转为list（.tolist()），丢失梯度追踪能力
• 使用math.log、abs等原生函数，而非TensorFlow的向量化操作
• 用random.shuffle打乱索引，无法被计算图捕获
• 循环遍历list处理anchor，而非TensorFlow向量化运算

修复方案：全TensorFlow化实现 #

重构损失函数 #

class RPN_Loss(tf.keras.losses.Loss):
    def L_cls(self, objectness_pred, objectness_true):
        # 数值稳定性处理，避免log(0)或log(1)报错
        objectness_pred = tf.clip_by_value(objectness_pred, 1e-7, 1 - 1e-7)
        cls_loss = -objectness_true * tf.math.log(objectness_pred) - (1 - objectness_true) * tf.math.log(1 - objectness_pred)
        return cls_loss

    def L1_smooth(self, d):
        # 向量化实现平滑L1损失，替代循环
        abs_d = tf.abs(d)
        return tf.where(abs_d < 1, 0.5 * tf.square(abs_d), abs_d - 0.5)

    def bb_loss(self, difference):
        # 对张量直接求和，无需遍历
        return tf.reduce_sum(self.L1_smooth(difference), axis=-1)

    def call(self, y_true, y_pred):
        batch_size = tf.shape(y_true)[0]
        image_width = tf.shape(y_true)[1]
        image_height = tf.shape(y_true)[2]
        k = tf.cast(tf.shape(y_true)[3] // 6, tf.int32)
        num_anchors = image_width * image_height * k

        # 保持Tensor格式，禁止转list
        y_true = tf.reshape(y_true, [-1, num_anchors, 6], order='F')
        y_pred = tf.reshape(y_pred, [-1, num_anchors, 6], order='F')

        mini_batch_size = 256
        l = tf.cast(1 / mini_batch_size * num_anchors, tf.float32)

        # 拆分标签与预测值
        offset_true = y_true[..., :4]
        objectness_true = y_true[..., 4:6]
        offset_pred = y_pred[..., :4]
        objectness_pred = y_pred[..., 4:6]

        # 生成正、负样本mask
        pos_mask = tf.reduce_all(tf.equal(objectness_true, [1, 0]), axis=-1)
        neg_mask = tf.reduce_all(tf.equal(objectness_true, [0, 1]), axis=-1)

        batch_losses = []
        for idx in range(batch_size):
            # 提取单样本的正负样本索引
            pos_indices = tf.where(pos_mask[idx])[:, 0]
            neg_indices = tf.where(neg_mask[idx])[:, 0]

            # 随机采样正样本（最多128个）
            pos_sample_num = tf.minimum(tf.shape(pos_indices)[0], mini_batch_size // 2)
            pos_indices = tf.random.shuffle(pos_indices)[:pos_sample_num]

            # 采样负样本凑够256个mini-batch
            neg_sample_num = mini_batch_size - pos_sample_num
            neg_indices = tf.random.shuffle(neg_indices)[:neg_sample_num]

            selected_indices = tf.concat([pos_indices, neg_indices], axis=0)

            # 提取选中样本的预测与真实值
            selected_obj_pred = tf.gather(objectness_pred[idx, :, 0], selected_indices)
            selected_obj_true = tf.gather(objectness_true[idx, :, 0], selected_indices)
            selected_offset_pred = tf.gather(offset_pred[idx], selected_indices)
            selected_offset_true = tf.gather(offset_true[idx], selected_indices)
            selected_pos_mask = tf.gather(pos_mask[idx], selected_indices)

            # 计算分类损失
            cls_loss = tf.reduce_sum(self.L_cls(selected_obj_pred, selected_obj_true))

            # 计算回归损失（仅正样本参与）
            offset_diff = selected_offset_pred - selected_offset_true
            reg_loss = tf.reduce_sum(tf.where(selected_pos_mask, self.bb_loss(offset_diff), 0.0))

            # 单样本损失计算
            sample_loss = cls_loss / tf.cast(mini_batch_size, tf.float32) + reg_loss * l / tf.cast(num_anchors, tf.float32)
            batch_losses.append(sample_loss)

        # 批次平均损失
        total_loss = tf.reduce_mean(tf.stack(batch_losses))
        return total_loss

模型结构微调说明 #

• objectness分支用softmax符合RPN二分类场景（标签为one-hot格式[1,0]/[0,1]），逻辑正确
• offset分支用tanh合理，因为RPN偏移量通常被归一化到[-1,1]区间

损失函数规范验证 #

你的实现完全符合Faster R-CNN原论文的RPN损失规范：

1. 分类损失：标准二元交叉熵，与原论文一致
2. 回归损失：平滑L1损失，避免L2损失在大误差下的梯度爆炸问题，符合原论文设计
3. 样本采样：固定mini-batch为256，正负样本比例1:1（最多128个正样本），严格遵循原论文策略
4. 损失加权：分类损失除以mini-batch大小，回归损失按正样本数量加权，逻辑与原论文对齐

额外注意事项 #

• 确保combined_output_train中忽略样本（既非正也非负）的标签为[0,0,0,0,0,0]，避免被错误计入损失
• run_eagerly=True仅用于调试，梯度问题修复后建议关闭以提升训练效率

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