Specify Custom Layer Backward Function

If Deep Learning Toolbox™ does not provide the layer you require for your task, then you can define your own custom layer. For a list of built-in layers, see List of Deep Learning Layers.

The example Define Custom Deep Learning Layer with Learnable Parameters shows how to create a custom SReLU layer and goes through the following steps:

Name the layer — Give the layer a name so that you can use it in MATLAB^®.
Declare the layer properties — Specify the properties of the layer, including learnable parameters and state parameters.
Create the constructor function (optional) — Specify how to construct the layer and initialize its properties. If you do not specify a constructor function, then at creation, the software initializes the Name, Description, and Type properties with [] and sets the number of layer inputs and outputs to 1.
Create initialize function (optional) — Specify how to initialize the learnable and state parameters when the software initializes the network. If you do not specify an initialize function, then the software does not initialize parameters when it initializes the network.
Create forward functions — Specify how data passes forward through the layer (forward propagation) at prediction time and at training time.
Create reset state function (optional) — Specify how to reset state parameters.
Create a backward function (optional) — Specify the derivatives of the loss with respect to the input data and the learnable parameters (backward propagation). If you do not specify a backward function, then the forward functions must support dlarray objects.

If the forward function only uses functions that support dlarray objects, then creating a backward function is optional. In this case, the software determines the derivatives automatically using automatic differentiation. For a list of functions that support dlarray objects, see List of Functions with dlarray Support. If you want to use functions that do not support dlarray objects, or want to use a specific algorithm for the backward function, then you can define a custom backward function using this example as a guide.

Create Custom Layer

The example Define Custom Deep Learning Layer with Learnable Parameters shows how to create a SReLU layer. A SReLU layer performs a thresholding operation, where for each channel, the layer scales values outside an interval. The interval thresholds and scaling factors are learnable parameters.[1]

The SReLU operation is given by

$f (x_{i}) = {\begin{array}{l} t_{i}^{l} + a_{i}^{l} (x_{i} - t_{i}^{l}) & if x_{i} \leq t_{i}^{l} \\ x_{i} & if t_{i}^{l} < x_{i} < t_{i}^{r} \\ t_{i}^{r} + a_{i}^{r} (x_{i} - t_{i}^{r}) & if t_{i}^{r} \leq x_{i} \end{array}$

where x_i is the input on channel i, t^l_i and t^r_i are the left and right thresholds on channel i, respectively, and a^l_i and a^r_i are the left and right scaling factors on channel i, respectively. These threshold values and scaling factors are learnable parameter, which the layer learns during training.

View the layer created in the example Define Custom Deep Learning Layer with Learnable Parameters. This layer does not have a backward function.

classdef sreluLayer < nnet.layer.Layer ...
        & nnet.layer.Acceleratable
    % Example custom SReLU layer.

    properties (Learnable)
    % Layer learnable parameters
    
        LeftSlope
        RightSlope
        LeftThreshold
        RightThreshold
    end

    methods
        function layer = sreluLayer(args) 
            % layer = sreluLayer creates a SReLU layer.
            %
            % layer = sreluLayer(Name=name) also specifies the
            % layer name.
    
            arguments
                args.Name = "";
            end
    
            % Set layer name.
            layer.Name = args.Name;

            % Set layer description.
            layer.Description = "SReLU";
        end

        function layer = initialize(layer,layout)
            % layer = initialize(layer,layout) initializes the layer
            % learnable parameters using the specified input layout.

            % Find number of channels.
            idx = finddim(layout,"C");
            numChannels = layout.Size(idx);

            % Initialize empty learnable parameters.
            sz = ones(1,numel(layout.Size);
            sz(idx) = numChannels;
            
            if isempty(layer.LeftSlope)
                layer.LeftSlope = rand(sz);
            end
            
            if isempty(layer.RightSlope)
                layer.RightSlope = rand(sz);
            end
            
            if isempty(layer.LeftThreshold)
                layer.LeftThreshold = rand(sz);
            end
            
            if isempty(layer.RightThreshold)
                layer.RightThreshold = rand(sz);
            end
        end

        function Y = predict(layer, X)
            % Y = predict(layer, X) forwards the input data X through the
            % layer and outputs the result Y.
            
            tl = layer.LeftThreshold;
            al = layer.LeftSlope;
            tr = layer.RightThreshold;
            ar = layer.RightSlope;
            
            Y = (X <= tl) .* (tl + al.*(X-tl)) ...
                + ((tl < X) & (X < tr)) .* X ...
                + (tr <= X) .* (tr + ar.*(X-tr));
        end
    end
end

Create Backward Function

Implement the backward function that returns the derivatives of the loss with respect to the input data and the learnable parameters.

The backward function syntax depends on the type of layer.

dLdX = backward(layer,X,Y,dLdY,memory) returns the derivatives dLdX of the loss with respect to the layer input, where layer has a single input and a single output. Y corresponds to the forward function output and dLdY corresponds to the derivative of the loss with respect to Y. The function input memory corresponds to the memory output of the forward function.
[dLdX,dLdW] = backward(layer,X,Y,dLdY,memory) also returns the derivative dLdW of the loss with respect to the learnable parameter, where layer has a single learnable parameter.
[dLdX,dLdSin] = backward(layer,X,Y,dLdY,dLdSout,memory) also returns the derivative dLdSin of the loss with respect to the state input, where layer has a single state parameter and dLdSout corresponds to the derivative of the loss with respect to the layer state output.
[dLdX,dLdW,dLdSin] = backward(layer,X,Y,dLdY,dLdSout,memory) also returns the derivative dLdW of the loss with respect to the learnable parameter and returns the derivative dLdSin of the loss with respect to the layer state input, where layer has a single state parameter and single learnable parameter.

You can adjust the syntaxes for layers with multiple inputs, multiple outputs, multiple learnable parameters, or multiple state parameters:

For layers with multiple inputs, replace X and dLdX with X1,...,XN and dLdX1,...,dLdXN, respectively, where N is the number of inputs.
For layers with multiple outputs, replace Y and dLdY with Y1,...,YM and dLdY1,...,dLdYM, respectively, where M is the number of outputs.
For layers with multiple learnable parameters, replace dLdW with dLdW1,...,dLdWP, where P is the number of learnable parameters.
For layers with multiple state parameters, replace dLdSin and dLdSout with dLdSin1,...,dLdSinK and dLdSout1,...,dLdSoutK, respectively, where K is the number of state parameters.

To reduce memory usage by preventing unused variables being saved between the forward and backward pass, replace the corresponding input arguments with ~.

Tip

If the number of inputs to backward can vary, then use varargin instead of the input arguments after layer. In this case, varargin is a cell array of the inputs, where the first N elements correspond to the N layer inputs, the next M elements correspond to the M layer outputs, the next M elements correspond to the derivatives of the loss with respect to the M layer outputs, the next K elements correspond to the K derivatives of the loss with respect to the K state outputs, and the last element corresponds to memory.

If the number of outputs can vary, then use varargout instead of the output arguments. In this case, varargout is a cell array of the outputs, where the first N elements correspond to the N the derivatives of the loss with respect to the N layer inputs, the next P elements correspond to the derivatives of the loss with respect to the P learnable parameters, and the next K elements correspond to the derivatives of the loss with respect to the K state inputs.

Because a SReLU layer has only one input, one output, four learnable parameters, and does not require the outputs of the layer forward function or a memory value, the syntax for backward for a SReLU layer is [dLdX,dLdLS,dLdRS,dldLT,dldRT] = backward(layer,X,~,dLdY,~). The dimensions of X are the same as in the forward function. The dimensions of dLdY are the same as the dimensions of the output Y of the forward function. The dimensions and data type of dLdX are the same as the dimensions and data type of X. The dimension and data type of dLdLS, dLdRS, dldLT, dldRT are the same as the dimension and data type of the learnable parameters LeftSlope, RightSlope, LeftThreshold, and RightThreshold, respectively.

During the backward pass, the layer automatically updates the learnable parameters using the corresponding derivatives.

To include a custom layer in a network, the layer forward functions must accept the outputs of the previous layer and forward propagate arrays with the size expected by the next layer. Similarly, when the layer specifies a backward function, it must accept inputs with the same size as the corresponding output of the forward function and backward propagate derivatives with the same size.

The derivative of the loss with respect to the input data is

$\frac{\partial L}{\partial x_{i}} = \frac{\partial L}{\partial f (x_{i})} \frac{\partial f (x_{i})}{\partial x_{i}}$

where $\partial L / \partial f (x_{i})$ is the gradient propagated from the next layer, and the derivative of the activation is

$\frac{\partial f (x_{i})}{\partial x_{i}} = {\begin{array}{l} a_{i}^{l} & if x_{i} \leq t_{i}^{l} \\ 1 & if t_{i}^{l} < x_{i} < t_{i}^{r} \\ a_{i}^{r} & if t_{i}^{r} \leq x_{i} \end{array} .$

The derivative of the loss with respect to the learnable parameter t^l_i is

$\frac{\partial L}{\partial t_{i}^{l}} = \sum_{j}^{} \frac{\partial L}{\partial f (x_{i j})} \frac{\partial f (x_{i j})}{\partial t_{j}^{l}}$

where i indexes the channels, j indexes the remaining elements, and the gradient of the activation is

$\frac{\partial f (x_{i})}{\partial t_{i}^{l}} = {\begin{array}{l} 1 - a_{i}^{l} & if x_{i} \leq t_{i}^{l} \\ 0 & if t_{i}^{l} < x_{i} < t_{i}^{r} \\ 0 & if t_{i}^{r} \leq x_{i} \end{array} .$

Similarly, for the other learnable parameters, the gradients are:

$\begin{array}{l} \frac{\partial f (x_{i})}{\partial t_{i}^{r}} = {\begin{array}{l} 0 & if x_{i} \leq t_{i}^{l} \\ 0 & if t_{i}^{l} < x_{i} < t_{i}^{r} \\ 1 - a_{i}^{r} & if t_{i}^{r} \leq x_{i} \end{array} \\ \frac{\partial f (x_{i})}{\partial a_{i}^{l}} = {\begin{array}{l} x_{i} - t_{i}^{l} & if x_{i} \leq t_{i}^{l} \\ 0 & if t_{i}^{l} < x_{i} < t_{i}^{r} \\ 0 & if t_{i}^{r} \leq x_{i} \end{array} \\ \frac{\partial f (x_{i})}{\partial a_{i}^{r}} = {\begin{array}{l} 0 & if x_{i} \leq t_{i}^{l} \\ 0 & if t_{i}^{l} < x_{i} < t_{i}^{r} \\ x_{i} - t_{i}^{r} & if t_{i}^{r} \leq x_{i} \end{array} \end{array}$

Create the backward function that returns these derivatives.

        function [dLdX,dLdLS,dLdRS,dLdLT,dLdRT] = backward(layer,X,~,dLdY,~)
            % [dLdX,dLdLS,dLdRS,dLdLT,dLdRT] = backward(layer,X,~,dLdY,~)
            % backward propagates the derivative of the loss function
            % through the layer.
            % Inputs:
            %         layer - Layer to backward propagate through
            %         X     - Input data
            %         dLdY  - Gradient propagated from the deeper layer
            % Outputs:
            %         dLdX  - Derivative of the loss with respect to the
            %                 input data
            %         dLdLS - Derivative of the loss with respect to the
            %                 learnable parameter LeftScale
            %         dLdRS - Derivative of the loss with respect to the
            %                 learnable parameter RightScale
            %         dLdLT - Derivative of the loss with respect to the
            %                 learnable parameter LeftThreshold
            %         dLdRT - Derivative of the loss with respect to the
            %                 learnable parameter RightThreshold

            ndims = numel(dims(X));
            idxC = finddim(X,"C");

            X = stripdims(X);
            dLdY = stripdims(dLdY);

            tl = layer.LeftThreshold;
            al = layer.LeftSlope;
            tr = layer.RightThreshold;
            ar = layer.RightSlope;

            dYdX = (X <= tl) .* al ...
                + (X > tl  & X < tr) ...
                + (X >= tr) .* ar;
            
            idx = setdiff(1:ndims,idxC);
            dLdLT = sum(dLdY .* (X <= tl) .* (1 - al),idx);
            dLdRT = sum(dLdY .* (tr <= X) .* (1 - ar),idx);
            dLdLS = sum(dLdY .* (X <= tl) .* (X - tl),idx);
            dLdRS = sum(dLdY .* (tr <= X) .* (X - tr),idx);
        end

Complete Layer

View the completed layer class file.

classdef sreluLayer < nnet.layer.Layer
    % Example custom SReLU layer.

    properties (Learnable)
    % Layer learnable parameters
    
        LeftSlope
        RightSlope
        LeftThreshold
        RightThreshold
    end
    
    methods
        function layer = sreluLayer(args) 
            % layer = sreluLayer creates a SReLU layer.
            %
            % layer = sreluLayer(Name=name) also specifies the
            % layer name.
    
            arguments
                args.Name = "";
            end
    
            % Set layer name.
            layer.Name = args.Name;

            % Set layer description.
            layer.Description = "SReLU";
        end

        function layer = initialize(layer,layout)
            % layer = initialize(layer,layout) initializes the layer
            % learnable parameters using the specified input layout.

            % Find number of channels.
            idx = finddim(layout,"C");
            numChannels = layout.Size(idx);

            % Initialize empty learnable parameters.
            sz = ones(1,numel(layout.Size);
            sz(idx) = numChannels;
            
            if isempty(layer.LeftSlope)
                layer.LeftSlope = rand(sz);
            end
            
            if isempty(layer.RightSlope)
                layer.RightSlope = rand(sz);
            end
            
            if isempty(layer.LeftThreshold)
                layer.LeftThreshold = rand(sz);
            end
            
            if isempty(layer.RightThreshold)
                layer.RightThreshold = rand(sz);
            end
        end

        function Y = predict(layer, X)
            % Y = predict(layer, X) forwards the input data X through the
            % layer and outputs the result Y.
            
            tl = layer.LeftThreshold;
            al = layer.LeftSlope;
            tr = layer.RightThreshold;
            ar = layer.RightSlope;
            
            Y = (X <= tl) .* (tl + al.*(X-tl)) ...
                + ((tl < X) & (X < tr)) .* X ...
                + (tr <= X) .* (tr + ar.*(X-tr));
        end
        
        function [dLdX,dLdLS,dLdRS,dLdLT,dLdRT] = backward(layer,X,~,dLdY,~)
            % [dLdX,dLdLS,dLdRS,dLdLT,dLdRT] = backward(layer,X,~,dLdY,~)
            % backward propagates the derivative of the loss function
            % through the layer.
            % Inputs:
            %         layer - Layer to backward propagate through
            %         X     - Input data
            %         dLdY  - Gradient propagated from the deeper layer
            % Outputs:
            %         dLdX  - Derivative of the loss with respect to the
            %                 input data
            %         dLdLS - Derivative of the loss with respect to the
            %                 learnable parameter LeftScale
            %         dLdRS - Derivative of the loss with respect to the
            %                 learnable parameter RightScale
            %         dLdLT - Derivative of the loss with respect to the
            %                 learnable parameter LeftThreshold
            %         dLdRT - Derivative of the loss with respect to the
            %                 learnable parameter RightThreshold

            ndims = numel(dims(X));
            idxC = finddim(X,"C");

            X = stripdims(X);
            dLdY = stripdims(dLdY);

            tl = layer.LeftThreshold;
            al = layer.LeftSlope;
            tr = layer.RightThreshold;
            ar = layer.RightSlope;

            dYdX = (X <= tl) .* al ...
                + (X > tl  & X < tr) ...
                + (X >= tr) .* ar;
            
            idx = setdiff(1:ndims,idxC);
            dLdLT = sum(dLdY .* (X <= tl) .* (1 - al),idx);
            dLdRT = sum(dLdY .* (tr <= X) .* (1 - ar),idx);
            dLdLS = sum(dLdY .* (X <= tl) .* (X - tl),idx);
            dLdRS = sum(dLdY .* (tr <= X) .* (X - tr),idx);
        end
    end
end

GPU Compatibility

If the layer forward functions fully support dlarray objects, then the layer is GPU compatible. Otherwise, to be GPU compatible, the layer functions must support inputs and return outputs of type gpuArray (Parallel Computing Toolbox).

Many MATLAB built-in functions support gpuArray (Parallel Computing Toolbox) and dlarray input arguments. For a list of functions that support dlarray objects, see List of Functions with dlarray Support. For a list of functions that execute on a GPU, see Run MATLAB Functions on a GPU (Parallel Computing Toolbox). To use a GPU for deep learning, you must also have a supported GPU device. For information on supported devices, see GPU Computing Requirements (Parallel Computing Toolbox). For more information on working with GPUs in MATLAB, see GPU Computing in MATLAB (Parallel Computing Toolbox).

References

[1] Hu, Xiaobin, Peifeng Niu, Jianmei Wang, and Xinxin Zhang. “A Dynamic Rectified Linear Activation Units.” IEEE Access 7 (2019): 180409–16. https://doi.org/10.1109/ACCESS.2019.2959036.