Custom loss function for Pixel Wise Cross Entropy Loss

autograd
1

I want to implement a custom loss function that looks like the one in the image below:

Capture
Capture1542×319 55.1 KB

I have never implemented a custom loss before so I am kinda stuck.

2

Hi,

If this is just the cross entropy loss for each pixel independently, then you can use the existing cross entropy provided by pytorch.
The pytorch function only accepts input of size (batch_dim, n_classes).
So if your output is of size (batch, height, width, n_classes), you can use .view(batch * height * width, n_classes) before giving it to the cross entropy function (considering each pixel as a different batch element) to achieve what you want.

2 Likes
3

Hey! Thank you for responding! My output is of the shape [100, 11, 80, 80] where batch = 100, n_classes = 11, height/width = 80. How do I convert this tensor to a tensor that will be accepted by cross entropy loss?

4

You can do out = inp.permute(0, 2, 3, 1).contiguous().view(-1, inp.size(1)) to have a Tensor whose shape is [640000, 11]