Model Trained for Image segmentation generates extremely similar outputs for different images

I am building my own image segmentation model to segment buildings from aerial images. This is how a batch of my dataset looks like:

I trained the model and I got the following losses:


I felt that the model will perform well. I tried giving it inputs from a test dataset. Bu the model generated similar outputs irrespective of the input. This is how the output looks like:

This is my model architecture:

class BuildingSegmenterNet(nn.Module):
    def __init__(self):
        super(BuildingSegmenterNet, self).__init__()
        self.seq1 = nn.Sequential(
            nn.Conv2d(3, 16, (5,5)),
        self.seq2 = nn.Sequential(
            nn.Conv2d(16, 8, (5,5)),
        self.seq3 = nn.Sequential(
            nn.Linear((61*61*8), 512),
        self.dropout1 = nn.Dropout(0.33)
        self.seq4 = nn.Sequential(
            nn.Linear(512, 128),
        self.dropout2 = nn.Dropout(0.33)
        self.seq5 = nn.Sequential(
            nn.Linear(128, 128),
        self.dropout3 = nn.Dropout(0.33)
        self.seq6 = nn.Sequential(
            nn.Linear(128, 512),
        self.dropout4 = nn.Dropout(0.33)
        self.seq7 = nn.Sequential(
            nn.Linear(512, 256*256),
    def forward(self, x):
        x = self.seq1(x)
        x = self.seq2(x)
        x = x.view(-1, 61*61*8)
        x = self.seq3(x)
        x = self.dropout1(x)
        x = self.seq4(x)
        x = self.dropout2(x)
        x = self.seq5(x)
        x = self.dropout3(x)
        x = self.seq6(x)
        x = self.dropout4(x)
        x = self.seq7(x)
        x = x.view(-1, 1, 256, 256)
        return x

And this is my loss function:

class SegmentationCrossEntropyLoss(nn.Module):
    def __init__(self):
        super(SegmentationCrossEntropyLoss, self).__init__()
    def forward(self, preds, targets):
        loss = nn.functional.binary_cross_entropy_with_logits(preds, targets.float())
        return loss

Edit: Add model arch and loss

1 Like