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