Semantic segmentation CNN outputs all zeroes

Question

I'm using MATLAB 2019, Linux, and UNet (a CNN specifically designed for semantic segmentation). I'm training the network to classify all pixels in an image as either cell or background to get segmentations of cells in microscopic images. My problem is the network is classifying every single pixel as background, and seems to just be outputting all zeroes. The validation accuracy improves a little at the very start of the training but than plateaus at around 60% for the majority of the training time. The network doesn't seem to be training very well and I have no idea why.

Can anyone give me some hints about what I should look into more closely? I just don't even know where to start with debugging this.

Here's my code:

    % Set datapath
    datapath = '/scratch/qbi/uqhrile1/ethans_lab_data';
    
    % Get training and testing datasets
    images_dataset = imageDatastore(strcat(datapath,'/bounding_box_cropped_resized_rgb'));
    load(strcat(datapath,'/gTruth.mat'));
    labels = pixelLabelDatastore(gTruth);
    [imdsTrain, imdsVal, imdsTest, pxdsTrain, pxdsVal, pxdsTest] = partitionCamVidData(images_dataset,labels);
    
    % Weight segmentation class importance by the number of pixels in each class
    pixel_count = countEachLabel(labels); % count number of each type of pixel
    frequency = pixel_count.PixelCount ./ pixel_count.ImagePixelCount; % calculate pixel type frequencies
    class_weights = mean(frequency) ./ frequency; % create class weights that balance the loss function so that more common pixel types won't be preferred
    
    % Specify the input image size.
    imageSize = [512 512 3];
    
    % Specify the number of classes.
    numClasses = 2;
    
    % Create network
    lgraph = unetLayers(imageSize,numClasses);
    
    % Replace the network's classification layer with a pixel classification
    % layer that uses class weights to balance the loss function
    pxLayer = pixelClassificationLayer('Name','labels','Classes',pixel_count.Name,'ClassWeights',class_weights);
    lgraph = replaceLayer(lgraph,"Segmentation-Layer",pxLayer);
    
    %% TRAIN THE NEURAL NETWORK
    
    % Define validation dataset-with-labels
    validation_dataset_with_labels = pixelLabelImageDatastore(imdsVal,pxdsVal);
    
    % Training hyper-parameters: edit these settings to fine-tune the network
    options = trainingOptions('adam', 'LearnRateSchedule','piecewise', 'LearnRateDropPeriod',10, 'LearnRateDropFactor',0.3, 'InitialLearnRate',1e-3, 'L2Regularization',0.005, 'ValidationData',validation_dataset_with_labels, 'ValidationFrequency',10, 'MaxEpochs',3, 'MiniBatchSize',1, 'Shuffle','every-epoch');
    
    % Set up data augmentation to enhance training dataset
    aug_imgs = {};
    numberOfImages = length(imdsTrain.Files);
    for k = 1 : numberOfImages
        % Apply cutout augmentation
        img = readimage(imdsTrain,k);
        cutout_img = random_cutout(img);imwrite(cutout_img,strcat('/scratch/qbi/uqhrile1/ethans_lab_data/augmented_dataset/img_',int2str(k),'.tiff'));
    end
    aug_imdsTrain = imageDatastore('/scratch/qbi/uqhrile1/ethans_lab_data/augmented_dataset');
    % Add other augmentations
    augmenter = imageDataAugmenter('RandXReflection',true, 'RandXTranslation',[-10 10],'RandYTranslation',[-10 10]);
    % Combine augmented data with training data
    augmented_training_dataset = pixelLabelImageDatastore(aug_imdsTrain, pxdsTrain, 'DataAugmentation',augmenter);
    
    % Train the network
    [cell_segmentation_nn, info] = trainNetwork(augmented_training_dataset,lgraph,options);
    
    save cell_segmentation_nn

Answer 1

Similar to other answers, I don't know Matlab that well but you could try the following steps to debug your problem.

Make sure you can overfit to a single instance

from your dataset, pull out a single image with a good amount of true positives in it. Duplicate that images B times (where B = Batch Size) and then try to train your network with only that small dataset. If you can't overfit to a single instance, then something is really wrong and you should validate all functional aspects of your network. If you can overfit, then it's probably more of an algorithmic or data imbalance issue.

Validate Functional Aspects of the Network

Make sure your input data and labels are correct.

Validate that your images are being correclty inputted into the network. You can do this by printing out any images right before they go into the training function. For the labels, make sure to manually inspect a few labels to be sure that the labels correctly match up the images.

Make sure that your loss function is correct.

Add a unit test or two to your loss function to validate that it is doing what it should be doing. Create a simiple example that you can easily validate.

Validate the rest of the non-algorithmic functionality

Anything that isn't a design choice should be validated. Make sure that the weights are the correct sizes, each layer has the correct number of weights, the intermediate features have the correct shape, etc.

Data Imbalance

If you were able to overfit to a single image, and validated the functional aspects of your network then you may be facing a data imbalance issue. Look at your dataset and see what % of your instances are true positives vs. true negatives. If you have an exteme imbalance (like 10% 90%) or something like that, then build a dataset that is more balanced and see if you can fit your data. If you fit the data with that more balanced dataset, then there's plenty of ways to fix your data imbalance issue. Google around for data imbalance and you should get a few good ideas. Some include focal loss, upsampling, etc...

Check your receptive field

The receptive field on a network is basically the area that can be looked at by the network for a specific region. This is controlled by the size of the convolution filters, stride, etc. If the data that you are segmenting doesn't fit within the receptive field of the model, then you might be saturating the loss function. tldr try playing around with kernel sizes.

Answer 2

I have never used MATLAB for ml before, so it is difficult for me to understand all your code. My first association to your problem is class imabalance. Since you seem to have got a handle on that, the problem could be dying ReLU or bloated activations. To check if the ReLU is dying, you could look at the activations of the early layers of your network. If many values are zero, it should be dying ReLU.