Custom cross-entropy loss in pytorch

Question

I have done a custom implementation of the pytorch cross-entropy loss function (as I need more flexibility to be introduced later). The model I intend to train with this will need a considerable amount of time to train and the resources available can't be used to merely test if the function is correct implementation. I have implemented vectorized implementation as it will be quicker to run.

Following is my code for the same:

def custom_cross(my_pred,true,batch_size=BATCH_SIZE):
    loss= -torch.mean(torch.sum(true.view(batch_size, -1) * torch.log(my_pred.view(batch_size, -1)), dim=1))
    return loss

I will really appreciate if you can suggest a more optimized implementation of the same or if I am making a mistake in the present one. The model will use a Nvidia Tesla K-80 to train.

Answer 1

If you need just cross entropy you can take the advantage PyTorch defined that.

import torch.nn.functional as F
loss_func = F.cross_entropy

suggest a more optimized implementation

PyTorch has F. loss functions, but you can easily write your own using plain python. PyTorch will create fast GPU or vectorized CPU code for your function automatically.

So, you may check the PyTorch original implementation but I think is this:

def log_softmax(x):
    return x - x.exp().sum(-1).log().unsqueeze(-1)

And here is the original implementation of cross entropy loss, now you may just alter:

nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)

To something you need, and you have it.