Currently I’m using as you suggested nn.BCEWithLogitsLoss and tried the model on a small data set of 15 images (should be small enough). But now the loss even if it goes down very fast it saturates and around 0.086 and 0.084.
So I guess this indicates that there is a more serious problem on the model or on the data?
I’ve tried out of curiosity this example using ResNetUnet from this example ResNetUnet and again the loss saturates for a small training data around 0.083.
I’ve also normalized my data with
transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
Any suggestion is more than welcome. Thanks again !