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

question about mutil-classes #15

Closed lbhnoteatfish closed 4 years ago

lbhnoteatfish commented 4 years ago

hi! I want to train in my cus dataset! It has 12+1(background) classes . I changed the line299 to softmax and change the line of model_test = model_unet(model_Inputs[0], 3, 12) . But if has a question that :::::: ValueError: Target size (torch.Size([4, 1, 96, 96])) must be the same as input size (torch.Size([4, 12, 96, 96])) I think it's the last layer of the network that outputs a feature map for each class. But my label is single channel ,it is shape is [4,1,96,96] . How can T calculate the loss ? Thank u !!!!

bigmb commented 4 years ago

Which line it is gving the error? Can you tell me more details.