Final layer of Unet for multiple classes?

Hi Guys,

I am trying to build Unet for one of my project from scratch but my label consists Multiple Classes(10 Classes).I gone through multiple source but unable to understand the final layer of Unet Network(Last layer of Decorder/Upscaling part in Unet ).

Currently I am using Linear Transformation in the end I don’t know it is a correct solution or not ??

Initial Layer of Down scaling

self.conv2d_1 = nn.Conv2d(1,32,kernel_size = 3,padding =1,stride=1)

Final Layer of Up scaling.

self.conv2d_10 = nn.Conv2d(32,32,kernel_size = 3,padding =1,stride=1)
self.Fc1 = nn.Linear(32,10)

In UNet, the output layer is a conv2d layer with kernel size 1 followed by a sigmoid activation

out_layer = nn.Conv2d(in_ch, num_classes, 1)

1 Like