Add CSpaceToDepthLayer and CDepthToSpaceLayer (#169)

* Add SpaceToDepth and DepthToSpace functions to MathEngine Signed-off-by: Valeriy Fedyunin <valery.fedyunin@abbyy.com> * Fix Metal kernels Signed-off-by: Valeriy Fedyunin <valery.fedyunin@abbyy.com> * Remove unused variable from CPU SpaceToDepthFunc Signed-off-by: Valeriy Fedyunin <valery.fedyunin@abbyy.com> * Delete unused test parts Signed-off-by: Valeriy Fedyunin <valery.fedyunin@abbyy.com> * Add CSpaceToDepthLayer and CDepthToSpaceLayer Signed-off-by: Valeriy Fedyunin <valery.fedyunin@abbyy.com> * Update version of NeoMLTest Signed-off-by: Valeriy Fedyunin <valery.fedyunin@abbyy.com> * Clarify docs Signed-off-by: Valeriy Fedyunin <valery.fedyunin@abbyy.com> * Fix checks in CDepthToSpaceLayer and CSpaceToDepthLayer Signed-off-by: Valeriy Fedyunin <valery.fedyunin@abbyy.com> * Add SpaceToDepth and DepthToSpace layers to Python library Signed-off-by: Valeriy Fedyunin <valery.fedyunin@abbyy.com> Co-authored-by: Stanislav Angeliuk <59917951+SAngeliuk@users.noreply.github.com>
neoml-lib · Jun 18, 2021 · 22395ea · 22395ea
1 parent 6b1802f
commit 22395ea
Show file tree

Hide file tree

Showing 42 changed files with 1,820 additions and 0 deletions.
diff --git a/NeoML/Python/CMakeLists.txt b/NeoML/Python/CMakeLists.txt
@@ -85,6 +85,7 @@ pybind11_add_module(PythonWrapper
     src/PySinkLayer.cpp
     src/PySolver.cpp
     src/PySourceLayer.cpp
+    src/PySpaceAndDepthLayer.cpp
     src/PySubSequenceLayer.cpp
     src/PyTiedEmbeddingsLayer.cpp
     src/PyTrainingModel.cpp

diff --git a/NeoML/Python/neoml/Dnn/SpaceAndDepth.py b/NeoML/Python/neoml/Dnn/SpaceAndDepth.py
@@ -0,0 +1,159 @@
+""" Copyright (c) 2017-2020 ABBYY Production LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+--------------------------------------------------------------------------------------------------------------*/
+"""
+
+import neoml.PythonWrapper as PythonWrapper
+from .Dnn import Layer
+from neoml.Utils import check_input_layers
+import neoml.Blob as Blob
+
+
+class SpaceToDepth(Layer):
+    """The layer that splits images into square blocks of size `k x k x Ch` and
+    writes the contents of these blocks to the corresponding pixels (`1 x 1 x Ch*k*k`) of the
+    output images in channel-last ordering.
+
+    As a result image of size `H x W x Ch` is transformed into images of size `H/k x W/k x Ch*k*k`.
+
+    This operation is the inverse function of DepthToSpace.
+
+    :param input_layers: The input layers to be connected. 
+        The integer in each tuple specifies the number of the output.
+        If not set, the first output will be used.
+    :type input_layers: object, tuple(object, int) or list of them
+    :param block_size: The size of the block (`k` from the formula).
+    :type block_size: int, default=1
+    :param name: The layer name.
+    :type name: str, default=None
+
+    .. rubric:: Layer inputs:
+
+    Has single input, of the dimensions:
+
+    - **BatchLength** * **BatchWidth** * **ListSize** is equal to the number of images
+    - **Height** is the image height; should be a multiple of `block_size`
+    - **Width** is the image width; should be a multiple of `block_size`
+    - **Depth** is equal to `1`
+    - **Channels** is the number of channels in the image format
+
+    .. rubric:: Layer outputs:
+
+    The layer has single output, of the dimensions:
+
+    - **BatchLength** is equal to the input `BatchLength`
+    - **BatchWidth** is equal to the input `BatchWidth`
+    - **ListSize** is equal to the input `ListSize`
+    - **Height** is equal to the input `Height / block_size`
+    - **Width** is equal to the input `Width / block_size`
+    - **Depth** is equal to `1`
+    - **Channels** is equal to the input `Channels * block_size * block_size`
+    """
+
+    def __init__(self, input_layer, block_size=1, name=None):
+        if type(input_layer) is PythonWrapper.SpaceToDepth:
+            super().__init__(input_layer)
+            return
+
+        layers, outputs = check_input_layers(input_layer, 1)
+
+        if block_size < 1:
+            raise ValueError('`block_size` must be < 0.')
+
+        internal = PythonWrapper.SpaceToDepth(str(name), layers[0], int(outputs[0]), block_size)
+        super().__init__(internal)
+
+    @property
+    def block_size(self):
+        """Gets the block size.
+        """
+        return self._internal.get_block_size()
+
+    @block_size.setter
+    def block_size(self, new_block_size):
+        """Sets the block size.
+        """
+        if new_block_size < 1:
+            raise ValueError('`block_size` must be < 0.')
+        self._internal.set_block_size(int(new_block_size))
+
+# ----------------------------------------------------------------------------------------------------------------------
+
+
+class DepthToSpace(Layer):
+    """The layer that transforms each pixel (`1 x 1 x Ch`) of 2-dimensional
+    images into square blocks of size `k x k x Ch/(k*k)`.
+    The elements of pixel are interpreted as an image of size `k x k x Ch/(k*k)` in channel-last ordering.
+    As a result `H x W x Ch` image is transformed into `H*k x W*k x Ch/(k*k)` image.
+
+    This operation is the inverse function of SpaceToDepth.
+
+    :param input_layers: The input layers to be connected. 
+        The integer in each tuple specifies the number of the output.
+        If not set, the first output will be used.
+    :type input_layers: object, tuple(object, int) or list of them
+    :param block_size: The size of the block (`k` from the formula).
+    :type block_size: int, default=1
+    :param name: The layer name.
+    :type name: str, default=None
+
+    .. rubric:: Layer inputs:
+
+    Has single input, of the dimensions:
+
+    - **BatchLength** * **BatchWidth** * **ListSize** is equal to the number of images
+    - **Height** is the image height
+    - **Width** is the image width
+    - **Depth** is equal to `1`
+    - **Channels** is the number of channels in the image format; should be a multiple of `block_size * block_size`
+
+    .. rubric:: Layer outputs:
+
+    The layer has single output, of the dimensions:
+
+    - **BatchLength** is equal to the input `BatchLength`
+    - **BatchWidth** is equal to the input `BatchWidth`
+    - **ListSize** is equal to the input `ListSize`
+    - **Height** is equal to the input `Height * block_size`
+    - **Width** is equal to the input `Width * block_size`
+    - **Depth** is equal to `1`
+    - **Channels** is equal to the input `Channels / ( block_size * block_size )`
+    """
+
+    def __init__(self, input_layer, block_size=1, name=None):
+        if type(input_layer) is PythonWrapper.DepthToSpace:
+            super().__init__(input_layer)
+            return
+
+        layers, outputs = check_input_layers(input_layer, 1)
+
+        if block_size < 1:
+            raise ValueError('`block_size` must be < 0.')
+
+        internal = PythonWrapper.DepthToSpace(str(name), layers[0], int(outputs[0]), block_size)
+        super().__init__(internal)
+
+    @property
+    def block_size(self):
+        """Gets the block size.
+        """
+        return self._internal.get_block_size()
+
+    @block_size.setter
+    def block_size(self, new_block_size):
+        """Sets the block size.
+        """
+        if new_block_size < 1:
+            raise ValueError('`block_size` must be < 0.')
+        self._internal.set_block_size(int(new_block_size))
diff --git a/NeoML/Python/neoml/Dnn/__init__.py b/NeoML/Python/neoml/Dnn/__init__.py
@@ -38,6 +38,7 @@
 from .Sink import Sink
 from .Softmax import Softmax
 from .Source import Source
+from .SpaceAndDepth import DepthToSpace, SpaceToDepth
 from .Split import SplitChannels, SplitDepth, SplitWidth, SplitHeight, SplitBatchWidth  
 from .SubSequence import SubSequence, ReverseSequence
 from .TiedEmbeddings import TiedEmbeddings

diff --git a/NeoML/Python/src/PySpaceAndDepthLayer.cpp b/NeoML/Python/src/PySpaceAndDepthLayer.cpp
@@ -0,0 +1,96 @@
+/* Copyright © 2017-2021 ABBYY Production LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+	http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+--------------------------------------------------------------------------------------------------------------*/
+
+#include <common.h>
+#pragma hdrstop
+
+#include "PySpaceAndDepthLayer.h"
+
+class CPySpaceToDepthLayer : public CPyLayer {
+public:
+	explicit CPySpaceToDepthLayer( CSpaceToDepthLayer& layer, CPyMathEngineOwner& mathEngineOwner ) : CPyLayer( layer, mathEngineOwner ) {}
+
+	void SetBlockSize( int value ) { Layer<CSpaceToDepthLayer>()->SetBlockSize( value ); }
+	int GetBlockSize() const { return Layer<CSpaceToDepthLayer>()->GetBlockSize(); }
+
+	py::object CreatePythonObject() const
+	{
+		py::object pyModule = py::module::import( "neoml.Dnn" );
+		py::object pyConstructor = pyModule.attr( "SpaceToDepth" );
+		return pyConstructor( py::cast( this ) );
+	}
+};
+
+class CPyDepthToSpaceLayer : public CPyLayer {
+public:
+	explicit CPyDepthToSpaceLayer( CDepthToSpaceLayer& layer, CPyMathEngineOwner& mathEngineOwner ) : CPyLayer( layer, mathEngineOwner ) {}
+
+	void SetBlockSize( int value ) { Layer<CDepthToSpaceLayer>()->SetBlockSize( value ); }
+	int GetBlockSize() const { return Layer<CDepthToSpaceLayer>()->GetBlockSize(); }
+
+	py::object CreatePythonObject() const
+	{
+		py::object pyModule = py::module::import( "neoml.Dnn" );
+		py::object pyConstructor = pyModule.attr( "DepthToSpace" );
+		return pyConstructor( py::cast( this ) );
+	}
+};
+
+void InitializeSpaceAndDepthLayer( py::module& m )
+{
+	py::class_<CPySpaceToDepthLayer, CPyLayer>( m, "SpaceToDepth" )
+		.def( py::init([]( const CPyLayer& layer )
+		{
+			return new CPySpaceToDepthLayer( *layer.Layer<CSpaceToDepthLayer>(), layer.MathEngineOwner() );
+		}))
+		.def( py::init([]( const std::string& name, const CPyLayer& inputLayer, int outputNumber, int blockSize )
+		{
+			CDnn& dnn = inputLayer.Dnn();
+			IMathEngine& mathEngine = dnn.GetMathEngine();
+
+			CPtr<CSpaceToDepthLayer> spaceToDepth = new CSpaceToDepthLayer( mathEngine );
+			spaceToDepth->SetName( FindFreeLayerName( dnn, "SpaceToDepth", name ).c_str() );
+			dnn.AddLayer( *spaceToDepth );
+			spaceToDepth->Connect( 0, inputLayer.BaseLayer(), outputNumber );
+			CPySpaceToDepthLayer* result = new CPySpaceToDepthLayer( *spaceToDepth, inputLayer.MathEngineOwner() );
+			result->SetBlockSize( blockSize );
+			return result;
+		}))
+		.def( "set_block_size", &CPySpaceToDepthLayer::SetBlockSize, py::return_value_policy::reference )
+		.def( "get_block_size", &CPySpaceToDepthLayer::GetBlockSize, py::return_value_policy::reference )
+	;
+
+	py::class_<CPyDepthToSpaceLayer, CPyLayer>( m, "DepthToSpace" )
+		.def( py::init([]( const CPyLayer& layer )
+		{
+			return new CPyDepthToSpaceLayer( *layer.Layer<CDepthToSpaceLayer>(), layer.MathEngineOwner() );
+		}))
+		.def( py::init([]( const std::string& name, const CPyLayer& inputLayer, int outputNumber, int blockSize )
+		{
+			CDnn& dnn = inputLayer.Dnn();
+			IMathEngine& mathEngine = dnn.GetMathEngine();
+
+			CPtr<CDepthToSpaceLayer> depthToSpace = new CDepthToSpaceLayer( mathEngine );
+			depthToSpace->SetName( FindFreeLayerName( dnn, "DepthToSpace", name ).c_str() );
+			dnn.AddLayer( *depthToSpace );
+			depthToSpace->Connect( 0, inputLayer.BaseLayer(), outputNumber );
+			CPyDepthToSpaceLayer* result = new CPyDepthToSpaceLayer( *depthToSpace, inputLayer.MathEngineOwner() );
+			result->SetBlockSize( blockSize );
+			return result;
+		}))
+		.def( "set_block_size", &CPyDepthToSpaceLayer::SetBlockSize, py::return_value_policy::reference )
+		.def( "get_block_size", &CPyDepthToSpaceLayer::GetBlockSize, py::return_value_policy::reference )
+	;
+}
diff --git a/NeoML/Python/src/PySpaceAndDepthLayer.h b/NeoML/Python/src/PySpaceAndDepthLayer.h
@@ -0,0 +1,20 @@
+/* Copyright © 2017-2021 ABBYY Production LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+	http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+--------------------------------------------------------------------------------------------------------------*/
+
+#pragma once
+
+#include "PyLayer.h"
+
+void InitializeSpaceAndDepthLayer( py::module& m );
diff --git a/NeoML/Python/src/PyWrapper.cpp b/NeoML/Python/src/PyWrapper.cpp
@@ -55,6 +55,7 @@ limitations under the License.
 #include "PyRepeatSequenceLayer.h"
 #include "PySequenceSumLayer.h"
 #include "PySoftmaxLayer.h"
+#include "PySpaceAndDepthLayer.h"
 #include "PySplitLayer.h"
 #include "PySubSequenceLayer.h"
 #include "PyTransformLayer.h"
@@ -130,6 +131,7 @@ PYBIND11_MODULE(PythonWrapper, m) {
 	InitializeTiedEmbeddingsLayer( m );
 	InitializeUpsampling2DLayer( m );
 	InitializeSourceLayer( m );
+	InitializeSpaceAndDepthLayer( m );
 	InitializeSinkLayer( m );
 
 	InitializeSolver( m );

diff --git a/NeoML/Python/tests.py b/NeoML/Python/tests.py
@@ -1445,6 +1445,41 @@ def test_transposedconv3d(self):
         self.assertEqual(out1.shape, (9, 3, 14, 11, 9, 7))
         self.assertEqual(out2.shape, (9, 3, 14, 11, 9, 7))
 
+    def test_depthtospace(self):
+        math_engine = neoml.MathEngine.CpuMathEngine(1)
+        dnn = neoml.Dnn.Dnn(math_engine)
+        source = neoml.Dnn.Source(dnn, 'source')
+        depth_to_space = neoml.Dnn.DepthToSpace(source, block_size=3, name='depth_to_space')
+        sink = neoml.Dnn.Sink(depth_to_space, 'sink')
+
+        self.assertEqual(depth_to_space.name, 'depth_to_space')
+        self.assertEqual(depth_to_space.block_size, 3)
+        depth_to_space.block_size = 2
+        self.assertEqual(depth_to_space.block_size, 2)
+
+        input_blob = neoml.Blob.asblob(math_engine, np.ones((2, 3, 5, 4, 8, 12), dtype=np.float32), (2, 3, 5, 4, 8, 1, 12))
+        outputs = dnn.run({'source' : input_blob})
+        out = outputs['sink'].asarray()
+        self.assertEqual(out.shape, (2, 3, 5, 8, 16, 3))
+
+    def test_spacetodepth(self):
+        math_engine = neoml.MathEngine.CpuMathEngine(1)
+        dnn = neoml.Dnn.Dnn(math_engine)
+        source = neoml.Dnn.Source(dnn, 'source')
+        space_to_depth = neoml.Dnn.SpaceToDepth(source, block_size=3, name='space_to_depth')
+        sink = neoml.Dnn.Sink(space_to_depth, 'sink')
+
+        self.assertEqual(space_to_depth.name, 'space_to_depth')
+        self.assertEqual(space_to_depth.block_size, 3)
+        space_to_depth.block_size = 2
+        self.assertEqual(space_to_depth.block_size, 2)
+
+        input_blob = neoml.Blob.asblob(math_engine, np.ones((2, 3, 5, 4, 8, 12), dtype=np.float32), (2, 3, 5, 4, 8, 1, 12))
+        outputs = dnn.run({'source' : input_blob})
+        out = outputs['sink'].asarray()
+        self.assertEqual(out.shape, (2, 3, 5, 2, 4, 48))
+
+
 class PoolingTestCase(TestCase):
     def _test_pooling(self, layer, init_params={}, changed_params={},
                       input_shape=(2, 1, 2, 3, 5, 4, 2)):