Normalized Cross Entropy Loss Implementation Tensorflow/Keras

Question

I am trying to implement a normalized cross entropy loss as described in this publication

The math given is:

This paper provided a PyTorch implementation:

@mlconfig.register
class NormalizedCrossEntropy(torch.nn.Module):
    def __init__(self, num_classes, scale=1.0):
        super(NormalizedCrossEntropy, self).__init__()
        self.device = device
        self.num_classes = num_classes
        self.scale = scale
    def forward(self, pred, labels):
        pred = F.log_softmax(pred, dim=1)
        label_one_hot = torch.nn.functional.one_hot(labels, self.num_classes).float().to(self.device)
        nce = -1 * torch.sum(label_one_hot * pred, dim=1) / (- pred.sum(dim=1))
        return self.scale * nce.mean()

But I need this to be translated to tensorflow for my ongoing project. Can anyone help me implement this normalized crossentropy loss in tensorflow?

Answer 1

I think is just a matter of translating methods name:

# given y_pred as 1-hot and y-true the multiclass probabilities
def NCE(y_true, y_pred):
     num = - tf.math.reduce_sum(tf.multiply(y_true, y_pred), axis=1)
     denom = -tf.math.reduce_sum(y_pred, axis=1)
     return tf.reduce_mean(num / denom)
t = tf.constant([[1,0,0], [0,0,1]], dtype=tf.float64)
y = tf.constant([[0.3,0.6,0.1], [0.1,0.1,0.8]], dtype=tf.float64)
NCE(t,y)
# <tf.Tensor: shape=(), dtype=float64, numpy=0.55>

Just check if the resulting loss is the same since I've not tested it