focal loss in keras

from keras import backend as K
import tensorflow as tf

# import dill

def binary_focal_loss(gamma=2., alpha=.25):
    """
    Binary form of focal loss.
      FL(p_t) = -alpha * (1 - p_t)**gamma * log(p_t)
      where p = sigmoid(x), p_t = p or 1 - p depending on if the label is 1 or 0, respectively.
    References:
        https://arxiv.org/pdf/1708.02002.pdf
    Usage:
     model.compile(loss=[binary_focal_loss(alpha=.25, gamma=2)], metrics=["accuracy"], optimizer=adam)
    """
    def binary_focal_loss_fixed(y_true, y_pred):
        """
        :param y_true: A tensor of the same shape as `y_pred`
        :param y_pred:  A tensor resulting from a sigmoid
        :return: Output tensor.
        """
        pt_1 = tf.where(tf.equal(y_true, 1), y_pred, tf.ones_like(y_pred))
        pt_0 = tf.where(tf.equal(y_true, 0), y_pred, tf.zeros_like(y_pred))

        epsilon = K.epsilon()
        # clip to prevent NaN's and Inf's
        pt_1 = K.clip(pt_1, epsilon, 1. - epsilon)
        pt_0 = K.clip(pt_0, epsilon, 1. - epsilon)
        
        # 若返回值为sum，则相当于alpha=n_batchsize,若返回为mean，则相当于alpha=1/n_class
        return -K.sum(alpha * K.pow(1. - pt_1, gamma) * K.log(pt_1)) \
               -K.sum((1 - alpha) * K.pow(pt_0, gamma) * K.log(1. - pt_0))
    return binary_focal_loss_fixed

def multi_label_weight_fl(activation = 'tanh',gamma = 2.,w = 0.40):
    def multi_label_loss(y_true,y_pred):
        if len(y_true.shape.as_list())>2:
            y_true = tf.reshape(y_true,[-1,n_class])
        if len(y_pred.shape.as_list())>2:
            y_pred = tf.reshape(y_pred,[-1,n_class])
        def conds(y_true,y_pred,i,ml_loss):
            return tf.less(i,K.shape(y_true)[0])
        def bodys(y_true,y_pred,i,ml_loss):
            y_true_i = tf.slice(y_true,[i,0],[1,n_class])
            y_pred_i = tf.slice(y_pred,[i,0],[1,n_class])

            one, zero =  tf.constant(1, dtype=tf.float32),tf.constant(0, dtype=tf.float32)
            ones = tf.gather_nd(y_pred_i, tf.where(tf.equal(y_true_i, one))) #返回y_true[i]中实为1 的标签们对应的概率值
            zeros = tf.gather_nd(y_pred_i, tf.where(tf.equal(y_true_i, zero))) #返回y_true[i]中实为0 的标签们对应的概率值
            # 复制以上两个向量，以便两两相减求差值
            p_repeat = tf.tile(ones,tf.shape(zeros))
            n_repeat = tf.reshape(tf.tile(tf.reshape(zeros,[-1,1]),[1,tf.shape(ones)[0]]),[-1])
            p_n_pairs = tf.concat([tf.reshape(p_repeat,[-1,1]),tf.reshape(n_repeat,[-1,1])],1)
            p_kl = p_n_pairs[:,0]-p_n_pairs[:,1]

            if activation == 'tanh':
                alpha = w* tf.pow(1-p_repeat,gamma)+ (1-w)* tf.pow(1+n_repeat,gamma)
            elif activation == 'sigmoid':
                alpha = w* tf.pow(1-p_repeat,gamma)+ (1-w)* tf.pow(n_repeat,gamma)
            alpha = tf.pow(alpha,gamma)
            ml_loss_i = tf.multiply(tf.cast(1/(tf.shape(p_n_pairs)[0]),dtype=tf.float32),
                                    tf.reduce_sum(tf.multiply(alpha,tf.exp(-(p_kl)))))
            ml_loss_i = tf.cond(tf.logical_or(tf.is_inf(ml_loss_i), tf.is_nan(ml_loss_i)), 
                                              lambda : tf.constant(0.0), lambda : ml_loss_i)
            ml_loss_i = tf.reshape(ml_loss_i,shape=[-1,])    
            ml_loss = tf.concat([ml_loss,ml_loss_i],axis=0)
            ml_loss = tf.reshape(ml_loss,shape=[-1,])
            i = tf.add(i,1)
            return [y_true,y_pred,i,ml_loss]
        i = tf.constant(0,dtype=tf.int32)
        ml_loss = tf.constant([0,],dtype=tf.float32)
        y_true,y_pred,i,ml_loss  = tf.while_loop(conds,bodys,[y_true,y_pred,i,ml_loss],
                                shape_invariants=[y_true.get_shape(),y_pred.get_shape(),i.get_shape(),tf.TensorShape([None,])])
        ml_loss = tf.slice(ml_loss,[1],[-1]) #维度从[n,1]转成[n,]
        return ml_loss
    return multi_label_loss
    
def categorical_focal_loss(gamma=2., alpha=.25):
    """
    Softmax version of focal loss.
           m
      FL = ∑  -alpha * (1 - p_o,c)^gamma * y_o,c * log(p_o,c)
          c=1
      where m = number of classes, c = class and o = observation
    Parameters:
      alpha -- the same as weighing factor in balanced cross entropy
      gamma -- focusing parameter for modulating factor (1-p)
    Default value:
      gamma -- 2.0 as mentioned in the paper
      alpha -- 0.25 as mentioned in the paper
    References:
        Official paper: https://arxiv.org/pdf/1708.02002.pdf
        https://www.tensorflow.org/api_docs/python/tf/keras/backend/categorical_crossentropy
    Usage:
     model.compile(loss=[categorical_focal_loss(alpha=.25, gamma=2)], metrics=["accuracy"], optimizer=adam)
    """
    def categorical_focal_loss_fixed(y_true, y_pred):
        """
        :param y_true: A tensor of the same shape as `y_pred`
        :param y_pred: A tensor resulting from a softmax
        :return: Output tensor.
        """
        # Scale predictions so that the class probas of each sample sum to 1
        y_pred /= K.sum(y_pred, axis=-1, keepdims=True)
        # Clip the prediction value to prevent NaN's and Inf's
        epsilon = K.epsilon()
        y_pred = K.clip(y_pred, epsilon, 1. - epsilon)
        # Calculate Cross Entropy
        cross_entropy = -y_true * K.log(y_pred)
        # Calculate Focal Loss
        loss = alpha * K.pow(1 - y_pred, gamma) * cross_entropy
        # Sum the losses in mini_batch(or K.mean)
        return K.sum(loss, axis=1)

    return categorical_focal_loss_fixed