How can I change the softmax layer with a custom one in classification problems using trainNetwork?

I am using Convolutional Neural Networks for deep learning classification in MATLAB R2018b, and I would like to use a custom softmax layer instead of the default one.

I tried to build a custom softmax layer using the Intermediate Layer Template present in Define Custom Deep Learning Layers, but when I train the net with trainNetwork I get the following error:

Error using trainNetwork (line xxx)
Invalid network.

Caused by:
    Layer 'WeightedClass': Missing softmax layer. A classification layer must be preceded by a softmax layer.

Code

This is the code which defines the custom softmax layer:

    classdef mySoftmaxLayer < nnet.layer.Layer
    % Custom softmax layer.
    
    properties (Learnable)
        % Layer learnable parameters.
        
    end
    
    methods
        function layer = mySoftmaxLayer(name)
            % layer = mySoftmaxLayer(name) creates a layer
            % and specifies the layer name.
            
            % Set layer name.
            layer.Name = name;
            
            % Set layer description.
            layer.Description = "My softmax layer";
        end
        
        function Z = predict(layer, X)
            % Z = predict(layer, X) forwards the input data X through the
            % layer and outputs the result Z.
            
            Z = myFunctionSoftmax(X);
        end
        
        function dLdX = backward(layer, X, Z, dLdZ, memory)
            % dLdX = backward(layer, X, Z, dLdZ, memory)
            % backward propagates the derivative of the loss function
            % through the layer.
            %
            % Inputs:
            %         layer    - Layer to backward propagate through 
            %         X        - Input data 
            %         Z        - Output of layer forward function 
            %         dLdZ     - Gradient propagated from the deeper layer 
            %         memory   - Memory value which can be used in backward
            %                    propagation
            % Outputs:
            %         dLdX     - Derivative of the loss with respect to the
            %                    input data
           
            dLdX = myDerivativeSoftmax(X) .* dLdZ;

        end
    end
end

This is the code in which the network is defined and trained:

%% Define the CNN architecture and training options
layers = [
    imageInputLayer([height width channels],'Name','Input');
    
    convolution2DLayer([height winSize],numFilters,'Name','Conv');
    batchNormalizationLayer('Name','Batch');
    reluLayer('Name','Relu')
    
    fullyConnectedLayer(numClasses,'Name','FullyConn');
    mySoftmaxLayer('MySoftmax')
    classificationLayer];

options = trainingOptions(...
    'sgdm',...
    'MaxEpochs',100,...
    'MiniBatchSize',256,...
    'InitialLearnRate',0.001);

%% Train the CNN
net = trainNetwork(XTrain,YTrain,layers,options);

Answer 1

[Pruthvi Muppavarapu]

Hi Davide,

In order to use a softmax layer before classification layer, you could define the ClassificationLayer as a custom regression output layer. Feel free to refer to the following documentation to create a custom regression output layer:

https://www.mathworks.com/help/deeplearning/ug/define-custom-regression-output-layer.html

To clarify further, you would need to have both the custom softmax layer, defined with

https://www.mathworks.com/help/releases/R2018a/nnet/ug/define-custom-deep-learning-layer.html

And the classification layer defined as a custom regression output layer that behaves the same as the classification output layer but does not require a softmaxLayer.

Hope this helps.

Regards,

Pruthvi

[Z]

@Pruthvi Muppavarapu Since the labels are categorical, what numbers should they be set to for regression? 0 and 1? -1 and 1?