errror while output is multiclass

bigmb / Unet-Segmentation-Pytorch-Nest-of-Unets

Implementation of different kinds of Unet Models for Image Segmentation - Unet , RCNN-Unet, Attention Unet, RCNN-Attention Unet, Nested Unet

MIT License

1.87k stars 345 forks source link

errror while output is multiclass #9

Closed Monologuethl closed 5 years ago

Monologuethl commented 5 years ago

I change this line into model_test = model_unet(model_Inputs[0], 3, 5)

raise ValueError("Target size ({}) must be the same as input size ({})".format(target.size(), input.size()))
ValueError: Target size (torch.Size([4, 1, 256, 256])) must be the same as input size (torch.Size([4, 5, 256, 256]))

my train images is 3 channels,and label images is 1 channel ,and label values is 0,1,2,3,4,could you tell me how solve this problem?

bigmb commented 5 years ago

You will have to add softmax at end. This was designed for binary classification, so try adding softmax at the output and check. Check the output of torch summary.