Error using trainNetwork (line 191) Invalid training data. The output size ([1 1 2]) of the last layer does not match the response size ([512 512 2]).

I am working with hyperspectral data. My goal is to identify flowers in an image. My image dimensions are [512 512 204] and my mask dimensions are [512 512 2]. When I run the code I get the following error:

Error using trainNetwork (line 191)

Invalid training data. The output size ([1 1 2]) of the last layer does not match the response size

([512 512 2]).

Here is my code for building the network. I used a network of type 2D CNN and ResNet-50.

%--------------------------Building the network------------------------------------------------
% Define the input image size for 2D CNN
imageSize = [512, 512, 204]; % 2D input: height x width x channels
% Define number of classes
numClasses = 2; % Background and Flowers
layers2D = [
    imageInputLayer(imageSize, 'Normalization', 'none', 'Name', 'input_2d')
    convolution2dLayer(3, 256, 'Padding', 'same', 'Name', 'conv2d_1', 'WeightLearnRateFactor', 10, 'BiasLearnRateFactor', 20, ...
    'Weights', randn([3 3 204 256], 'single'), 'Bias', zeros(1, 1, 256, 'single')) % Bias with shape [1, 1, 256]
    batchNormalizationLayer('Name', 'batchnorm2d_1', ...
    'TrainedMean', zeros(1, 1, 256, 'single'), ...
    'TrainedVariance', ones(1, 1, 256, 'single'), ...
    'Offset', zeros(1, 1, 256, 'single'), ...
    'Scale', ones(1, 1, 256, 'single'))
    
    reluLayer('Name', 'relu2d_1')
    convolution2dLayer(3, 128, 'Padding', 'same', 'Name', 'conv2d_2', 'WeightLearnRateFactor', 10, 'BiasLearnRateFactor', 20, ...
    'Weights', randn([3 3 256 128], 'single'), 'Bias', zeros(1, 1, 128, 'single')) % Bias with shape [1, 1, 128]
    batchNormalizationLayer('Name', 'batchnorm2d_2', ...
    'TrainedMean', zeros(1, 1, 128, 'single'), ...
    'TrainedVariance', ones(1, 1, 128, 'single'), ...
    'Offset', zeros(1, 1, 128, 'single'), ...
    'Scale', ones(1, 1, 128, 'single'))
    
    reluLayer('Name', 'relu2d_2')
    maxPooling2dLayer(2, 'Stride', 2, 'Name', 'pool2d_1')
    convolution2dLayer(3, 64, 'Padding', 'same', 'Name', 'conv2d_3', 'WeightLearnRateFactor', 10, 'BiasLearnRateFactor', 20, ...
    'Weights', randn([3 3 128 64], 'single'), 'Bias', zeros(1, 1, 64, 'single')) % Bias with shape [1, 1, 64]
    batchNormalizationLayer('Name', 'batchnorm2d_3', ...
    'TrainedMean', zeros(1, 1, 64, 'single'), ...
    'TrainedVariance', ones(1, 1, 64, 'single'), ...
    'Offset', zeros(1, 1, 64, 'single'), ...
    'Scale', ones(1, 1, 64, 'single'))
    
    reluLayer('Name', 'relu2d_3')
    ];
% Create the layer graph
DeepLabv_2D = layerGraph(layers2D);
% % Visualize the network
% figure;
% plot(DeepLabv_2D);
% title('2D-CNN Network');
% analyzeNetwork(DeepLabv_2D)
% ---------------ResNet-50--------------------
% Load ResNet-50
net = resnet50;
ResNet50 = layerGraph(net);
% Modify the input layer to match hyperspectral data dimensions
newInputLayer = imageInputLayer([512 512 204], 'Name', 'input_1', 'Normalization', 'none');
ResNet50 = replaceLayer(ResNet50, 'input_1', newInputLayer);
% Modify the first convolutional layer to accept 204 channels instead of 3
oldConvLayer = ResNet50.Layers(2); % Get existing conv1 layer
newConvLayer = convolution2dLayer(oldConvLayer.FilterSize, oldConvLayer.NumFilters, ...
    'Stride', oldConvLayer.Stride, ...
    'Padding', oldConvLayer.PaddingSize, ...
    'NumChannels', 64, ...
    'WeightsInitializer', 'he', ...
    'BiasInitializer', 'zeros', ...
    'Name', 'conv1_custom'); 
newConvLayer.Weights = randn([oldConvLayer.FilterSize, 64, oldConvLayer.NumFilters]) * sqrt(2 / (64 * prod(oldConvLayer.FilterSize)));
newConvLayer.Bias = zeros(1, 1, oldConvLayer.NumFilters);
% Replace the original convolutional layer
ResNet50 = replaceLayer(ResNet50, 'conv1', newConvLayer);
% Replace the final layers for binary classification
newFullyConnectedLayer = fullyConnectedLayer(2, 'Name', 'fc_classification');
numInputs = 2048; 
newFullyConnectedLayer.Weights = randn(2, numInputs) * sqrt(2 / numInputs);
newFullyConnectedLayer.Bias = zeros(2, 1);
newSoftmaxLayer = softmaxLayer('Name', 'softmax');
categories = {'flowers', 'non_flowers'}; 
newClassificationLayer = classificationLayer('Name', 'classification_output', 'Classes', categories);
% Remove original classification layers
ResNet50 = removeLayers(ResNet50, {'fc1000', 'fc1000_softmax', 'ClassificationLayer_fc1000'});
% Add new classification layers
ResNet50 = addLayers(ResNet50, [newFullyConnectedLayer, newSoftmaxLayer, newClassificationLayer]);
% Connect layers appropriately
ResNet50 = connectLayers(ResNet50, 'avg_pool', 'fc_classification');
% % Visualize the modified network
% figure;
% plot(ResNet50);
% title('ResNet-50 Network');
% analyzeNetwork(ResNet50)
% ---------------------Connecting the networks-------------------------
% Now connect DeepLabv3 output to ResNet50
ResNet50 = addLayers(ResNet50, DeepLabv_2D.Layers);
ResNet50 = disconnectLayers(ResNet50, 'input_1', 'conv1_custom');
ResNet50 = removeLayers(ResNet50, 'input_1');
DeepLabResNet_Fusion_Network = connectLayers(ResNet50, 'relu2d_3', 'conv1_custom');
% Visualize the network
figure;
plot(DeepLabResNet_Fusion_Network);
title('The Updated Network');
analyzeNetwork(DeepLabResNet_Fusion_Network);
%----------------------------Network training-----------------------------------------
% Combine datastores for training
combinedData = combine(trainImages, trainMasks);
% Combine datastores for validation
validationData = combine(valImages, valMasks);
% Combine datastores for testing
testData = combine(testImages, testMasks);

% Check mask formats
% Check train masks
trainMasksExample = read(trainMasks);
disp('Train masks size:');
disp(size(trainMasksExample));
% Check validation masks
valMasksExample = read(valMasks);
disp('Validation masks size:');
disp(size(valMasksExample));
% Check test masks
testMasksExample = read(testMasks);
disp('Test masks size:');
disp(size(testMasksExample));
% Training options
options = trainingOptions('adam', ...
    'MaxEpochs', 50, ...
    'MiniBatchSize', 8, ...
    'InitialLearnRate', 1e-4, ...
    'Shuffle', 'every-epoch', ...
    'ValidationData', validationData, ...
    'ValidationFrequency', 5, ...
    'Verbose', true, ...
    'Plots', 'training-progress');
% Train the network
trainedNet = trainNetwork(combinedData, DeepLabResNet_Fusion_Network, options);

Answer 1

[Matt J]

It doesn't make sense that you have a "mask" unless you are doing semantic segmentation (i.e. pixel-by-pixel classsifcation), and if you are doing semantic segmentation, I don't think you will be able to repurpose ResNet, which is not normally used for that. You should probably follow the documentation's example on semantic segmentation,

https://www.mathworks.com/help/vision/ug/semantic-segmentation-using-deep-learning.html

In any case, no semantic segmentation networks uses a fully connected layer or classificationLayer in its final 3 layers. Those are used when you are trying to assign a single class label to the entire image.

[Hadas]

thanks

[Matt J]

You're weclome, but please Accept-click the answer if your question has been resolved.