How to set the input as a two-dimensional feature matrix in the custom Intermediate-layer of deep learning?

introduction:

As can be seen from the official documents, the custom Intermediate-layer has the following forms:

classdef myLayer < nnet.layer.Layer

    properties
        % (Optional) Layer properties.

        % Layer properties go here.
    end

    properties (Learnable)
        % (Optional) Layer learnable parameters.

        % Layer learnable parameters go here.
    end
    
    methods
        function layer = myLayer()
            % (Optional) Create a myLayer.
            % This function must have the same name as the class.

            % Layer constructor function goes here.
        end
        
        function [Z1, …, Zm] = predict(layer, X1, …, Xn)
            % Forward input data through the layer at prediction time and
            % output the result.
            %
            % Inputs:
            %         layer       - Layer to forward propagate through
            %         X1, ..., Xn - Input data
            % Outputs:
            %         Z1, ..., Zm - Outputs of layer forward function
            
            % Layer forward function for prediction goes here.
        end

        function [Z1, …, Zm, memory] = forward(layer, X1, …, Xn)
            % (Optional) Forward input data through the layer at training
            % time and output the result and a memory value.
            %
            % Inputs:
            %         layer       - Layer to forward propagate through
            %         X1, ..., Xn - Input data
            % Outputs:
            %         Z1, ..., Zm - Outputs of layer forward function
            %         memory      - Memory value for custom backward propagation

            % Layer forward function for training goes here.
        end

        function [dLdX1, …, dLdXn, dLdW1, …, dLdWk] = ...
                backward(layer, X1, …, Xn, Z1, …, Zm, dLdZ1, …, dLdZm, memory)
            % (Optional) Backward propagate the derivative of the loss  
            % function through the layer.
            %
            % Inputs:
            %         layer             - Layer to backward propagate through
            %         X1, ..., Xn       - Input data
            %         Z1, ..., Zm       - Outputs of layer forward function            
            %         dLdZ1, ..., dLdZm - Gradients propagated from the next layers
            %         memory            - Memory value from forward function
            % Outputs:
            %         dLdX1, ..., dLdXn - Derivatives of the loss with respect to the
            %                             inputs
            %         dLdW1, ..., dLdWk - Derivatives of the loss with respect to each
            %                             learnable parameter
            
            % Layer backward function goes here.
        end
    end
end

my issue:

when the network is forwarded, how to define the input parameter "X" of the "forward" or "predict" function as a two-dimensional input feature matrix, the label form is 'batchSize*numChannels' instead of ' height*wight*numChannels*batchSize' The default four-dimensional feature matrix?

specific problem：

For example, in the field of face recognition, I customize an "ArcFace Head loss" layer, the complete content is as follows, the syntax seems to be correct, but when I use the "checkLayer" function to verify the custom layer, there is no guarantee that "X" is two-dimensional Feature matrix?

classdef arcfaceLayer < nnet.layer.Layer
    % custom ArcFace Head loss 
    % reference：https://blog.csdn.net/yiran103/article/details/83684613
    % https://github.com/cuixing158/jetson_faceTrack_pytorch/blob/main/net/model.py
    % offical document: Define Custom Deep Learning Layer with Learnable Parameters
    %
    properties
        cos_m
        sin_m
        mm
        threshold
        s
        m
    end
    properties (Learnable)
        % Layer learnable parameters
            
        % kernel coefficient ,size: embedding_size*classnum
        kernel 
    end
    
    methods
        function layer = arcfaceLayer(name,embedding_size, classnum,  s, m) 
            % init function
            arguments
                name (1,1) string = "arcface_head"
                embedding_size (1,1) {mustBeNumeric}=512
                classnum (1,1) {mustBeNumeric}=4
                s (1,1) {mustBeNumeric}=64 % scalar value default is 64, see normface https://arxiv.org/abs/1704.06369
                m (1,1) {mustBeNumeric}=0.5 % the margin value, default is 0.5
            end
            % Set layer name.
            layer.Name = name;

            % Set layer description.
            layer.Description = "ArcFace with " + string(embedding_size) + " embedding size"+" classnum:"+string(classnum)+" scalar:"+string(s)+" m:"+string(m);
        
            % Initialize kernel coefficient
            layer.kernel = 2*rand(embedding_size,classnum)-1; 
            
            layer.cos_m = cos(m);
            layer.sin_m = sin(m);
            layer.mm = m.*sin(m);
            layer.threshold = cos(pi-m);
            layer.s = s;
            layer.m = m;
        end
        
        function Z = predict(layer, X)
            % Z = predict(layer, X) forwards the input data X through the
            % layer and outputs the result Z.
            %
            % 输入：
            %     X 为 bs*embedding_size二维矩阵
            % 输出：
            %     Z， 定义的大小为bs*classnum.
            %
            X = X./vecnorm(X,2,2);
            kernel_norm = layer.kernel./vecnorm(layer.kernel,2,1); % embedding_size*classnum
            
            % cos(theta+m)
            cos_theta = X*kernel_norm; % bs*classnum
            cos_theta_2 = cos_theta.^2;
            sin_theta_2 = 1 - cos_theta_2;
            sin_theta = sqrt(sin_theta_2);
            cos_theta_m = (cos_theta * layer.cos_m - sin_theta * layer.sin_m); %  bs*classnum,数学公式, cos(alpha+beta) = cos(alpha)*cos(beta)-sin(alpha)*sin(beta)
            % this condition controls the theta+m should in range [0, pi]
            %      0<=theta+m<=pi
            %   -m<=theta<=pi-m
            cond_v = cos_theta - layer.threshold; % bs*classnum, 按道理此值大于0
            cond_mask = cond_v <= 0; % bs*classnum, 找出小于0的索引
            keep_val = cos_theta - layer.mm; % when theta not in [0,pi], use cosface instead,# bs*classnum
            cos_theta_m(cond_mask) = keep_val(cond_mask); % bs*classnum

            output = cos_theta_m;
            Z = layer.s.*output; % scale up in order to make softmax work, first introduced in normface
        end
    end
end

Then ,check my layer:

layer = arcfaceLayer();
validInputSize = [10,512];
checkLayer(layer,validInputSize,'ObservationDimension',2) % or 'ObservationDimension' set to 3,4

failed???

______

Failure Summary:

Name Failed Incomplete Reason(s)

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

nnet.checklayer.TestLayerWithoutBackward/predictDoesNotError(Observations=one) X Failed by verification.

----------------------------------------------------------------------------------------------------------------------------------------------

nnet.checklayer.TestLayerWithoutBackward/predictDoesNotError(Observations=multiple) X Failed by verification.

----------------------------------------------------------------------------------------------------------------------------------------------

nnet.checklayer.TestLayerWithoutBackward/predictIsConsistentInType(Precision=single,Device=cpu) X Filtered by assumption.

----------------------------------------------------------------------------------------------------------------------------------------------

nnet.checklayer.TestLayerWithoutBackward/predictIsConsistentInType(Precision=single,Device=gpu) X Filtered by assumption.

----------------------------------------------------------------------------------------------------------------------------------------------

nnet.checklayer.TestLayerWithoutBackward/predictIsConsistentInType(Precision=double,Device=cpu) X Filtered by assumption.

----------------------------------------------------------------------------------------------------------------------------------------------

nnet.checklayer.TestLayerWithoutBackward/predictIsConsistentInType(Precision=double,Device=gpu) X Filtered by assumption.

----------------------------------------------------------------------------------------------------------------------------------------------

nnet.checklayer.TestLayerWithoutBackward/backwardPropagationDoesNotError(Observations=one) X Filtered by assumption.

----------------------------------------------------------------------------------------------------------------------------------------------

nnet.checklayer.TestLayerWithoutBackward/backwardPropagationDoesNotError(Observations=multiple) X Filtered by assumption.

Test Summary:

9 Passed, 2 Failed, 6 Incomplete, 4 Skipped.

Time elapsed: 14.8035 seconds.

Desired result：

• The custom layer does not report errors when checking

• Support two-dimensional feature matrix input

[cui,xingxing]

arcFace-matlab project is here, unspport layer