spatialToDepth custom layer problems

Hi all,

I tried to make a custom layer, following interesting article arXiv:1909.03205v2.

The layer re-arrange the input image such that all pixels remain, reducing the spatial dimensions and enlarging the number of feature maps, rather than use the distructive maxPooling.

The code for the layer is given below. Unfortunatelly, I got errors related to wrong output class from the layer and/or dimensions mismatch to the next layer which will be conv2D.

Please help, what should I do to make this layer passing checkLayer/analyzeNetwork?

BTW, running the predict function lines at the command window shows that the function is working well and do its job.

Kind regards,

Michael.

classdef spaceToDepthLayer < nnet.layer.Layer
    properties
        Stride;
        Padding;
    end
    methods
        function layer = spaceToDepthLayer(stride, padding, name)
            layer.Name = name;
            layer.Description = sprintf('Rearrangement spatial to depth: stride = [%d,%d]', stride(1), stride(2));
            layer.Stride = stride;      %2 element vector
            layer.Padding = padding;    %4 element vector [t b l r]
        end
        function Z = predict(layer, X)
            X = padarray(X, layer.Padding([1 3]), 'pre');
            X = padarray(X, layer.Padding([2 4]), 'post');
            [H, W, C] = size(X);
            Z = reshape(X, H, W*C);
            Z = Z';
            Z = reshape(Z, layer.Stride(2), W/layer.Stride(2), H);
            Z = permute(Z, [1 3 2]);
            Z = reshape(Z, layer.Stride(1)*layer.Stride(2), []);
            Z = reshape(Z', H/layer.Stride(1), W/layer.Stride(2), []);
        end
    end
end

[michael kaster]

Hi all,

It works very well in Python.