usuyama/pytorch-unet

network output

satrialoka opened this issue · 1 comments

Using this implementation for 2 class problem (I'm using 1 channel target image with binary(0,1) value), but my output is not binary (0,1). did I miss something? should I add some activation function at the last layer?

Yes, you need to apply sigmoid to the model output as in the example:

# Predict
pred = model(inputs)
# The loss functions include the sigmoid function.
pred = F.sigmoid(pred)
pred = pred.data.cpu().numpy()

The reason behind this is to use F.binary_cross_entropy_with_logits, which combines BCE+sigmoid in one layer and applies the log-sum-exp trick for numerical stability.