[SYSTEMDS-3694] Python NN Sequence and layer interface

This commit:
- Adds a Layer interface for the Python API.
- Affine and ReLU classes are changed to extend this interface.
- Fixes fixes some small formatting issues in the modified classes.
- Adds a Sequential primitive to the nn Python API.
  It is able to combine multiple nn layers into one sequential module.
- fix in the python MultiReturn so outputs of the instance can be properly accessed.
- Adds the backwards pass to the Sequential primitives.
- Variations to Sequential testing involving MultiReturns.
- Test if the input gradient is set correctly on the backwards pass
  and Fixes a bug where this was not the case on the affine layer.
- Testing to verify that the layer gets updated correctly during forward and backward pass.

AMLS project SoSe'24

Closes apache#2025

src/main/python/systemds/operator/nn/affine.py

@@ -18,33 +18,24 @@
 # under the License.
 #
 # -------------------------------------------------------------
-import os
-
 from systemds.context import SystemDSContext
-from systemds.operator import Matrix, Source, MultiReturn
-from systemds.utils.helpers import get_path_to_script_layers
+from systemds.operator import Matrix, MultiReturn
+from systemds.operator.nn.layer import Layer
 
 
-class Affine:
- _source: Source = None
+class Affine(Layer):
  weight: Matrix
  bias: Matrix
 
- def __new__(cls, *args, **kwargs):
- return super().__new__(cls)
-
  def __init__(self, sds_context: SystemDSContext, d, m, seed=-1):
  """
  sds_context: The systemdsContext to construct the layer inside of
  d: The number of features that are input to the affine layer
  m: The number of neurons that are contained in the layer, 
  and the number of features output
  """
- Affine._create_source(sds_context)
-
- # bypassing overload limitation in python
- self.forward = self._instance_forward
- self.backward = self._instance_backward
+ super().__init__(sds_context, 'affine.dml')
+ self._X = None
 
  # init weight and bias
  self.weight = Matrix(sds_context, '')
@@ -64,7 +55,7 @@ def forward(X: Matrix, W: Matrix, b: Matrix):
  b: The bias added in the output.
  return out: An output matrix.
  """
- Affine._create_source(X.sds_context)
+ Affine._create_source(X.sds_context, "affine.dml")
  return Affine._source.forward(X, W, b)
 
  @staticmethod
@@ -77,7 +68,7 @@ def backward(dout:Matrix, X: Matrix, W: Matrix, b: Matrix):
  return dX, dW, db: The gradients of: input X, weights and bias.
  """
  sds = X.sds_context
- Affine._create_source(sds)
+ Affine._create_source(sds, "affine.dml")
  params_dict = {'dout': dout, 'X': X, 'W': W, 'b': b}
  dX = Matrix(sds, '')
  dW = Matrix(sds, '')
@@ -104,11 +95,6 @@ def _instance_backward(self, dout: Matrix, X: Matrix):
  X: The input to this layer.
  return dX, dW,db: gradient of input, weights and bias, respectively
  """
- return Affine.backward(dout, X, self.weight, self.bias)
-
- @staticmethod
- def _create_source(sds: SystemDSContext):
- if Affine._source is None or Affine._source.sds_context != sds:
- path = get_path_to_script_layers()
- path = os.path.join(path, "affine.dml")
- Affine._source = sds.source(path, "affine")
+ gradients = Affine.backward(dout, X, self.weight, self.bias)
+ self._X = gradients[0]
+ return gradients

src/main/python/systemds/operator/nn/layer.py

@@ -0,0 +1,69 @@
+# -------------------------------------------------------------
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you 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 os
+
+from systemds.context import SystemDSContext
+from systemds.operator import Source
+from systemds.utils.helpers import get_path_to_script_layers
+
+
+class Layer:
+ """
+ Interface for neural network layers
+ """
+
+ _source: Source = None
+
+ def __init__(self, sds_context: SystemDSContext = None, dml_script: str = None):
+ if sds_context is not None and dml_script is not None:
+ self.__class__._create_source(sds_context, dml_script)
+
+ # bypassing overload limitation in python
+ self.forward = self._instance_forward
+ self.backward = self._instance_backward
+
+ @classmethod
+ def _create_source(cls, sds_context: SystemDSContext, dml_script: str):
+ """
+ Create SystemDS source
+ :param sds_context: SystemDS context
+ :param dml_script: DML script inside /scripts/nn/layers/
+ :return:
+ """
+ if cls._source is None or cls._source.sds_context != sds_context:
+ script_path = get_path_to_script_layers()
+ path = os.path.join(script_path, dml_script)
+ name = dml_script.split(".")[0]
+ cls._source = sds_context.source(path, name)
+
+ def _instance_forward(self, *args):
+ raise NotImplementedError
+
+ def _instance_backward(self, *args):
+ raise NotImplementedError
+
+ @staticmethod
+ def forward(*args):
+ raise NotImplementedError
+
+ @staticmethod
+ def backward(*args):
+ raise NotImplementedError

src/main/python/systemds/operator/nn/relu.py

@@ -18,28 +18,24 @@
 # under the License.
 #
 # -------------------------------------------------------------
-import os.path
-
 from systemds.context import SystemDSContext
 from systemds.operator import Matrix, Source
-from systemds.utils.helpers import get_path_to_script_layers
+from systemds.operator.nn.layer import Layer
 
 
-class ReLU:
+class ReLU(Layer):
  _source: Source = None
 
- def __init__(self, sds: SystemDSContext):
- ReLU._create_source(sds)
- self.forward = self._instance_forward
- self.backward = self._instance_backward
+ def __init__(self, sds_context: SystemDSContext):
+ super().__init__(sds_context, "relu.dml")
 
  @staticmethod
  def forward(X: Matrix):
  """
  X: input matrix
  return out: output matrix
  """
- ReLU._create_source(X.sds_context)
+ ReLU._create_source(X.sds_context, "relu.dml")
  return ReLU._source.forward(X)
 
  @staticmethod
@@ -49,7 +45,7 @@ def backward(dout: Matrix, X: Matrix):
  X: input matrix
  return dX: gradient of input
  """
- ReLU._create_source(dout.sds_context)
+ ReLU._create_source(dout.sds_context, "relu.dml")
  return ReLU._source.backward(dout, X)
 
  def _instance_forward(self, X: Matrix):
@@ -58,11 +54,3 @@ def _instance_forward(self, X: Matrix):
 
  def _instance_backward(self, dout: Matrix, X: Matrix):
  return ReLU.backward(dout, X)
-
- @staticmethod
- def _create_source(sds: SystemDSContext):
- if ReLU._source is None or ReLU._source.sds_context != sds:
- path = get_path_to_script_layers()
- path = os.path.join(path, "relu.dml")
- ReLU._source = sds.source(path, "relu")
-

src/main/python/systemds/operator/nn/sequential.py

@@ -0,0 +1,97 @@
+# -------------------------------------------------------------
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you 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.
+#
+# -------------------------------------------------------------
+from systemds.operator import MultiReturn
+from systemds.operator.nn.layer import Layer
+
+
+class Sequential(Layer):
+ def __init__(self, *args):
+ super().__init__()
+
+ self.layers = []
+ if len(args) == 1 and isinstance(args[0], list):
+ self.layers = args[0]
+ else:
+ self.layers = list(args)
+
+ def __len__(self):
+ return len(self.layers)
+
+ def __getitem__(self, idx):
+ return self.layers[idx]
+
+ def __setitem__(self, idx, value):
+ self.layers[idx] = value
+
+ def __delitem__(self, idx):
+ del self.layers[idx]
+
+ def __iter__(self):
+ return iter(self.layers)
+
+ def __reversed__(self):
+ return reversed(self.layers)
+
+ def push(self, layer: Layer):
+ """
+ Add layer
+ :param layer: Layer
+ :return:
+ """
+ self.layers.append(layer)
+
+ def pop(self):
+ """
+ Remove last layer
+ :return: Layer
+ """
+ return self.layers.pop()
+
+ def _instance_forward(self, X):
+ """
+ Forward pass
+ :param X: Input matrix
+ :return: output matrix
+ """
+ out = X
+ for layer in self:
+ out = layer.forward(out)
+
+ # if MultiReturn, take only output matrix
+ if isinstance(out, MultiReturn):
+ out = out[0]
+ return out
+
+ def _instance_backward(self, dout, X):
+ """
+ Backward pass
+ :param dout: gradient of output, passed from the upstream
+ :param X: input matrix
+ :return: output matrix
+ """
+ dx = dout
+ for layer in reversed(self):
+ dx = layer.backward(dx, X)
+
+ # if MultiReturn, take only gradient of input
+ if isinstance(dx, MultiReturn):
+ dx = dx[0]
+ return dx

src/main/python/systemds/operator/nodes/multi_return.py

@@ -47,7 +47,7 @@ def __init__(self, sds_context, operation,
  named_input_nodes, OutputType.MULTI_RETURN, False)
 
  def __getitem__(self, key):
- self._outputs[key]
+ return self._outputs[key]
 
  def code_line(self, var_name: str, unnamed_input_vars: Sequence[str],
  named_input_vars: Dict[str, str]) -> str:

src/main/python/tests/nn/test_affine.py

@@ -77,6 +77,7 @@ def test_forward(self):
  out = affine.forward(Xm).compute()
  self.assertEqual(len(out), 5)
  self.assertEqual(len(out[0]), 6)
+ assert_almost_equal(affine._X.compute(), Xm.compute())
 
  # test static method
  out = Affine.forward(Xm, Wm, bm).compute()
@@ -91,10 +92,13 @@ def test_backward(self):
 
  # test class method
  affine = Affine(self.sds, dim, m, 10)
- [dx, dw, db] = affine.backward(doutm, Xm).compute()
+ gradients = affine.backward(doutm, Xm)
+ intermediate = affine._X.compute()
+ [dx, dw, db] = gradients.compute()
  assert len(dx) == 5 and len(dx[0]) == 6
  assert len(dw) == 6 and len(dx[0]) == 6
  assert len(db) == 1 and len(db[0]) == 6
+ assert_almost_equal(intermediate, dx)
 
  # test static method
  res = Affine.backward(doutm, Xm, Wm, bm).compute()