smoothL1 loss的tensorflow实现

Posted 2022-09-14 MachineLP

MachineLP的Github（欢迎follow）：https://github.com/MachineLP

def smooth_l1_loss(bbox_pred, bbox_targets, bbox_inside_weights, bbox_outside_weights, sigma=1.0, dim=[1]):
    '''
    bbox_pred   ：预测框
    bbox_targets：标签框
    bbox_inside_weights：
    bbox_outside_weights：
    '''  
    sigma_2 = sigma ** 2
    box_diff = bbox_pred - bbox_targets
    in_box_diff = bbox_inside_weights * box_diff
    abs_in_box_diff = tf.abs(in_box_diff)
    # tf.less 返回 True or False； a<b,返回True， 否则返回False。
    smoothL1_sign = tf.stop_gradient(tf.to_float(tf.less(abs_in_box_diff, 1. / sigma_2)))
    # 实现公式中的条件分支
    in_loss_box = tf.pow(in_box_diff, 2) * (sigma_2 / 2.) * smoothL1_sign + (abs_in_box_diff - (0.5 / sigma_2)) * (1. - smoothL1_sign)
    out_loss_box = bbox_outside_weights * in_loss_box
    loss_box = tf.reduce_mean(tf.reduce_sum(out_loss_box, axis=dim))
    return loss_box