Back propagation on a deep learning custom layer with a maxpooling "alike" functionality

I'm trying to use the new deep learning package from matlab to try to define a custom layer for a project I'm working on. The general idea of the layer is to take as and output of such layer, the N values of neurons with the biggest energy given by the input. Very similar to a maxPooling layer but without having into account the spatial representation of this neurons. Since this layer will be at the end of certain convolutional blocks, a reshape is needed inside the layer. Ths idea is to go on with a fullyConnected layer after.

The code used as a guide: https://es.mathworks.com/help/nnet/ug/define-custom-deep-learning-layer.html

After many changes using the "checkLayer" function, this is the fardest I have gotten:

Code:

classdef globalmaxPooling < nnet.layer.Layer

    properties (Learnable)
        Alpha
    end

    methods
        function layer = globalmaxPooling(numChannels, name) 
        % Create an globalmaxPooling with numChannels as an input wich
        % represents the number of outputs this layer will have in a vector nature

            % Set layer name
            if nargin == 2
                layer.Name = name;
            end

            % Set layer description
            layer.Description = ...
                ['Performs a global max pooling of ', num2str(numChannels), ' features']; 

            % Initialize scaling coefficient
            layer.Alpha = numChannels;
        end

        function Z = predict(layer, X)
            N = size(X,4);
            X_reshaped = reshape(X,[size(X,1)*size(X,2)*size(X,3) N]);
            n_channels = layer.Alpha;
            Z = zeros([n_channels N]);
            if class(X_reshaped(:))== "gpuArray"
                Z = gpuArray(Z);
                for i = 1 : N
                    [~,I] = sort(X_reshaped(:,i),'descend');
                    Z(:,i) = gather(X_reshaped(gather(I(1:n_channels)),i));
                end
                Z = single(Z);
            elseif class(X_reshaped(:))== "double"
                Z = double(Z);
                for i = 1 : N
                    [~,I] = sort(X_reshaped(:,i),'descend');
                    Z(:,i) = X_reshaped(I(1:n_channels),i);
                end
                Z = double(Z);
            else
                Z = single(Z);
                for i = 1 : N
                    [~,I] = sort(X_reshaped(:,i),'descend');
                    Z(:,i) = X_reshaped(I(1:n_channels),i);
                end
                Z = single(Z);               
            end
        end

        function [Z, memory] = forward(layer, X) 

            N = size(X,4);
            X_reshaped = reshape(X,[size(X,1)*size(X,2)*size(X,3) N]);

            n_channels = layer.Alpha;
            Z = zeros([9216 N]);
            if class(X_reshaped(:))== "gpuArray"
                Z = gpuArray(Z);
                for i = 1 : N
                    [~,I] = sort(X_reshaped(:,i),'descend');
                    Z(:,i) = gather(X_reshaped(gather(flip(I(1:n_channels))),i));
                end
                Z = single(Z);
            elseif class(X_reshaped(:))== "double"
                Z = double(Z);
                for i = 1 : N
                    [~,I] = sort(X_reshaped(:,i),'descend');
                    Z(:,i) = X_reshaped(flip(I(1:n_channels)),i);
                end
                Z = double(Z);
            else
                Z = single(Z);
                for i = 1 : N
                    [~,I] = sort(X_reshaped(:,i),'descend');
                    Z(:,i) = X_reshaped(flip(I(1:n_channels)),i);
                end
                Z = single(Z);               
            end
            memory = [];
        end

        function [dLdX,dLdalpha] = backward(layer, X, Z, dLdZ, memory)      
            if class(X(:))== "gpuArray"

            end
            N = size(X,4);
            n_channels = layer.Alpha;
            X_reshaped = reshape(X,[size(X,1)*size(X,2)*size(X,3) N]);
            max_index = zeros(n_channels,N);
            for i = 1 : N
                [~,I] = sort(X_reshaped(:,i),'descend');
                max_index(:,i) = flip(gather(I(1:n_channels)));
            end           
            N = size(dLdZ,2);
            n_feat_out = size(X,1)*size(X,2)*size(X,3);
            dLdX_prima = zeros([n_feat_out N]);
            if class(X(:))== "gpuArray"
                dLdX_prima = gpuArray(dLdX_prima);
                for i = 1 : N
                    dLdX_prima(max_index(:,i),i) = dLdZ(:,i);
                end
                dLdX = single(reshape(dLdX_prima,[13 13 256 N]));

            elseif class(X(:))== "double"
                dLdX_prima = double(dLdX_prima);
                for i = 1 : N
                    dLdX_prima(max_index(:,i),i) = dLdZ(:,i);
                end
                dLdX = double(reshape(dLdX_prima,[13 13 256 N]));
            else
                dLdX_prima = single(dLdX_prima);
                for i = 1 : N
                    dLdX_prima(max_index(:,i),i) = dLdZ(:,i);
                end
                dLdX = single(reshape(dLdX_prima,[13 13 256 N]));
            end
            dLdalpha = layer.Alpha;
        end
    end
end

As a way to check the layer, I used the following command.

checkLayer(globalmaxPooling(9216,'globalmaxPooling5'),[13 13 256],'ObservationDimension',4)

But I got the following error:

================================================================================

Verification failed in nnet.checklayer.TestCase/gradientsAreNumericallyCorrect.

    ----------------
    Test Diagnostic:
    ----------------
    The derivative 'dLdX' for 'backward' is inconsistent with the numerical gradient. Either 'dLdX' is incorrectly computed, the function is non-differentiable at some input points, or the error tolerance is too small.

I have tried looking everywhere, but seems that there are not many people using this package, which makes it hard to find information about it.

Answer 1

[Joss Knight]

Don't write this:

        function [Z, memory] = forward(layer, X) 
            N = size(X,4);
            X_reshaped = reshape(X,[size(X,1)*size(X,2)*size(X,3) N]);
            n_channels = layer.Alpha;
            Z = zeros([9216 N]);
            if class(X_reshaped(:))== "gpuArray"
                Z = gpuArray(Z);
                for i = 1 : N
                    [~,I] = sort(X_reshaped(:,i),'descend');
                    Z(:,i) = gather(X_reshaped(gather(flip(I(1:n_channels))),i));
                end
                Z = single(Z);
            elseif class(X_reshaped(:))== "double"
                Z = double(Z);
                for i = 1 : N
                    [~,I] = sort(X_reshaped(:,i),'descend');
                    Z(:,i) = X_reshaped(flip(I(1:n_channels)),i);
                end
                Z = double(Z);
            else
                Z = single(Z);
                for i = 1 : N
                    [~,I] = sort(X_reshaped(:,i),'descend');
                    Z(:,i) = X_reshaped(flip(I(1:n_channels)),i);
                end
                Z = single(Z);               
            end
            memory = [];
        end

Instead, write this:

        function [Z, memory] = forward(layer, X) 
            N = size(X,4);
            X_reshaped = reshape(X, [], N);
            n_channels = layer.Alpha;
            Z = zeros([9216 N], 'like', X);

            for i = 1 : N
                [~,I] = sort(X_reshaped(:,i),'descend');
                Z(:,i) = X_reshaped(flip(I(1:n_channels)),i);
            end

            memory = [];
        end

Now, in answer to your question, what that bit of checkLayer is trying to do is tell you whether or not you have written your backward function correctly - i.e. have you differentiated your function correctly? If you are sure you have, for all the inputs you care about, then ignore the error and go ahead and use your layer.

[Alex garcia]

Regarding the correction you pointed out, this was the only way for the layerchecker didn't complain about the class of the output variables being different from the input class variables.

I believe that the function is correctly done, but when putting it inside a net, it gave me the following error:

Unexpected error calling cuDNN: CUDNN_STATUS_BAD_PARAM.

Caused by: Error using nnet.internal.cnngpu.softmaxBackward2D Unexpected error calling cuDNN: CUDNN_STATUS_BAD_PARAM.

This happened when I applied this layer to the well known alexNet, by just switching the last maxPool layer for my custom layer.

 layers = [
        alexnet.Layers(1:15)
        globalmaxPooling(9216,'global maxpooling')

        fullyConnectedLayer(4096,'Name','fc6')
        reluLayer('Name','relu6')
        dropoutLayer('Name','drop6')

        fullyConnectedLayer(4096,'Name','fc7')
        reluLayer('Name','relu7')
        dropoutLayer('Name','drop7')

        fullyConnectedLayer(100,'Name','fc8')
        softmaxLayer('Name','prob')
        classificationLayer('Name','output')
        ];

I read that this error is common in older matlab versions, but I am using matalb 2018a. I'm running a nVidia 930m in my laptop.

[Joss Knight]

But I've replaced your call to zeros with zeros(..., 'like', X), which will do the job of making sure Z is initialized with both the correct type and the correct storage medium (CPU or GPU). Those calls to gather are also superfluous. So your three separate blocks are all identical and just dealt with by constructing Z correctly. You've also being doing all your computation for GPU inputs on the CPU instead and in double precision, which seems inefficient.

Either way, something's going wrong because softmaxBackward is receiving input of the wrong type or size, which means your layer's forward method is outputting data of the wrong type or size. It looks on the face of it like you haven't reshaped Z back to H-by-W-by-C-by-N format, instead you've left it 2-D.

[Alex garcia]

I finally solved it. As u said there was something wrong with the size. For some reason, I couldn't output the correct matrix size out of my layer, trying it with a HxWxCxN and (H*W*C)xN. I added a standard 1x1 maxPooling2dLayer that made no changes to the its input, and it worked.

 layers = [
        alexnet.Layers(1:15)

        maxPooling2dLayer(1,'stride',[1 1],'Name','pool3')

        globalmaxPooling(9216,'global maxpooling')
        fullyConnectedLayer(4096,'Name','fc6')
        reluLayer('Name','relu6')
        dropoutLayer('Name','drop6')
        fullyConnectedLayer(4096,'Name','fc7')
        reluLayer('Name','relu7')
        dropoutLayer('Name','drop7')
        fullyConnectedLayer(100,'Name','fc8')
        softmaxLayer('Name','prob')
        classificationLayer('Name','output')
        ];

I believe that the network is working properly since, even though it is really slow, it returns some results.