Hello.
I’m trying to implement the Unet paper 1 from scratch but I’m struggling implementing what they call the “energy function”.
Citing the paper:

U-Net__Convolutional_Networks_for_Biomedical_Image_Segmentation.png1188×300 59.2 KB

So far this is what I have:
class EnergyFunction(nn.Module):
    def __init__(self):
        super(EnergyFunction, self).__init__()

    def forward(self, logits):
        # logits is of shape [batch_size, feature_channels, height, width]
        # Here: [4, 2, 388, 388]
        r = F.softmax(logits)
        # r is of shape: [4, 2, 388, 388]
        return r

But according to the paper I should pass in each of my 2 388x388 matrices (A and B) pixels to a softmax function where each resulting “pixel” is a vector of {[A(0,0), B(0, 0)], [A(0, 1), B(0, 1)]…} and so on. So basically I should have something like this: {softmax([A(0,0), B(0, 0)]), softmax([A(0, 1), B(0, 1)])…} right?
And the resulting shape of r should be [4, 1, 388, 388] right?
I really have no idea how to do that.
Thank you for your help.