Arcface loss, sphereface loss

I developed my own arcface layer which works well for image retrieval task. I used a pretrained ResNet101 and removed FC and loss layer. Then I added the my own FC layer and arcface loss layer to lgraph by using following script, where num_classes is the number of classes in training datat:

num_classes=20000;
in=2048;
out=num_classes;
sc = sqrt(3/(1*1*in)) ;
weights = (rand(out,in, 'single')*2 - 1)*sc ;
CFC = cosineFC('CFC',weights);
CFC=setLearnRateFactor(CFC,'Weights',10);
factor1=getLearnRateFactor(CFC,'Weights');
lgraph = addLayers(lgraph,CFC);
lgraph = connectLayers(lgraph,'gap','CFC');

The cosineFC layer computes the dot product between L2-normalized weights and L2-normalised features

classdef cosineFC < nnet.layer.Layer
    % Example custom PReLU layer.
    properties (Learnable)
        % Layer learnable parameters
            
        % Scaling coefficient
        Weights
    end
    
    methods
        function layer = cosineFC(name,Weights)
            layer.Name = name;
            % Set layer description.
            layer.Description = "cosineFC ";
        
            % Initialize scaling coefficient.
            layer.Weights = Weights; 
        end
        
        function Z = predict(layer, X)
            % Z = predict(layer, X) forwards the input data X through the
            X=L2_Norm(X,3);
            layer.Weights=L2_Norm(layer.Weights,2);
            bias = zeros(size(layer.Weights,1),1);
            Z  = fullyconnect(X,layer.Weights,bias,'DataFormat','SSCB'); 
            Z=permute(Z,[3 4 1 2]);
        end
    end

Finally I modified regression loss for arcface loss optimization

classdef arcloss < nnet.layer.RegressionLayer
    % Example custom regression layer with mean-absolute-error loss.
    
    methods
        function layer = arcloss(name)
            % layer = maeRegressionLayer(name) creates a
            % mean-absolute-error regression layer and specifies the layer
            % name.
			
            % Set layer name.
            layer.Name = name;
            % Set layer description.
            layer.Description = 'arcface loss';
        end
        
        function loss = forwardLoss(layer, cosine, Y)
            cosine=dlarray(squeeze(cosine),'CB');
            Y=squeeze(Y);
            m=0.3;
            s=30;
            cos_m = cos(m);
            sin_m = sin(m);
            th = cos(pi - m);
            mm = sin(pi - m) * m;
            cs=(1-cosine.^2);
            cs_limit=max(min(cs,1),0);
            sine=sqrt(cs_limit);
            phi = cosine * cos_m - sine * sin_m;
            cos_th=cosine-mm;
            idx=cosine<th;
            phi(idx) = cos_th(idx);
            output = (Y .* phi) + ((1 - Y) .* cosine);
            output = output*s;
            dlYPred = softmax(output);
            loss = crossentropy(dlYPred,Y);
        end
    end
end