feat: Implement Forward and Backward pass in Reshape Layer

Implementation of the Forward and Backward pass in the Reshape Layer,
which transforms the input to the desired output dimensions.

tmva/tmva/inc/TMVA/DNN/DeepNet.h

@@ -135,7 +135,7 @@ class TDeepNet {
    /*! Function for adding Reshape Layer in the Deep Neural Network, with a given
     *  height and width. It will take every matrix from the previous layer and
     *  reshape it to a matrix with new dimensions. */
-   TReshapeLayer<Architecture_t> *AddReshapeLayer(size_t height, size_t width);
+   TReshapeLayer<Architecture_t> *AddReshapeLayer(size_t depth, size_t height, size_t width);
 
    /*! Function fir adding Reshape Layer in the Deep Neural Network, when
     *  the layer is already created. */
@@ -444,13 +444,13 @@ void TDeepNet<Architecture_t, Layer_t>::AddDenseLayer(TDenseLayer<Architecture_t
 
 //______________________________________________________________________________
 template <typename Architecture_t, typename Layer_t>
-TReshapeLayer<Architecture_t> *TDeepNet<Architecture_t, Layer_t>::AddReshapeLayer(size_t height, size_t width)
+TReshapeLayer<Architecture_t> *TDeepNet<Architecture_t, Layer_t>::AddReshapeLayer(size_t depth, size_t height,
+                                                                                  size_t width)
 {
    size_t batchSize = this->GetBatchSize();
    size_t inputDepth;
    size_t inputHeight;
    size_t inputWidth;
-   size_t depth;
    size_t outputNSlices = this->GetBatchSize();
    size_t outputNRows;
    size_t outputNCols;
@@ -466,8 +466,7 @@ TReshapeLayer<Architecture_t> *TDeepNet<Architecture_t, Layer_t>::AddReshapeLaye
       inputWidth = lastLayer->GetWidth();
    }
 
-   depth = inputDepth;
-   outputNRows = inputDepth;
+   outputNRows = depth;
    outputNCols = height * width;
 
    TReshapeLayer<Architecture_t> *reshapeLayer = new TReshapeLayer<Architecture_t>(

tmva/tmva/inc/TMVA/DNN/ReshapeLayer.h

@@ -76,7 +76,11 @@ TReshapeLayer<Architecture_t>::TReshapeLayer(size_t batchSize, size_t inputDepth
    : VGeneralLayer<Architecture_t>(batchSize, inputDepth, inputHeight, inputWidth, depth, height, width, 0, 0, 0, 0,
                                    outputNSlices, outputNRows, outputNCols, EInitialization::kZero)
 {
-   // Nothing to do here.
+   if (this->GetInputDepth() * this->GetInputHeight() * this->GetInputWidth() !=
+       this->GetDepth() * this->GetHeight() * this->GetWidth()) {
+      std::cout << "Dimensions not compatible" << std::endl;
+      return;
+   }
 }
 
 //_________________________________________________________________________________________________
@@ -104,13 +108,19 @@ TReshapeLayer<Architecture_t>::~TReshapeLayer()
 template <typename Architecture_t>
 auto TReshapeLayer<Architecture_t>::Forward(std::vector<Matrix_t> input, bool applyDropout) -> void
 {
+   for (size_t i = 0; i < this->GetBatchSize(); i++) {
+      Architecture_t::Reshape(this->GetOutputAt(i), input[i]);
+   }
 }
 
 //_________________________________________________________________________________________________
 template <typename Architecture_t>
 auto TReshapeLayer<Architecture_t>::Backward(std::vector<Matrix_t> &gradients_backward,
                                              const std::vector<Matrix_t> &activations_backward) -> void
 {
+   for (size_t i = 0; i < this->GetBatchSize(); i++) {
+      Architecture_t::Reshape(gradients_backward[i], this->GetActivationGradientsAt(i));
+   }
 }
 
 //_________________________________________________________________________________________________

tmva/tmva/src/MethodDL.cxx

@@ -592,6 +592,7 @@ void MethodDL::ParseReshapeLayer(DNN::TDeepNet<Architecture_t, Layer_t> &deepNet
                                  std::vector<DNN::TDeepNet<Architecture_t, Layer_t>> &nets, TString layerString,
                                  TString delim)
 {
+   int depth = 0;
    int height = 0;
    int width = 0;
 
@@ -603,12 +604,16 @@ void MethodDL::ParseReshapeLayer(DNN::TDeepNet<Architecture_t, Layer_t> &deepNet
 
    for (; token != nullptr; token = (TObjString *)nextToken()) {
       switch (idxToken) {
-      case 1: // height
+      case 1: {
+         TString strDepth(token->GetString());
+         depth = strDepth.Atoi();
+      } break;
+      case 2: // height
       {
          TString strHeight(token->GetString());
          height = strHeight.Atoi();
       } break;
-      case 2: // width
+      case 3: // width
       {
          TString strWidth(token->GetString());
          width = strWidth.Atoi();
@@ -618,7 +623,7 @@ void MethodDL::ParseReshapeLayer(DNN::TDeepNet<Architecture_t, Layer_t> &deepNet
    }
 
    // Add the reshape layer
-   TReshapeLayer<Architecture_t> *reshapeLayer = deepNet.AddReshapeLayer(height, width);
+   TReshapeLayer<Architecture_t> *reshapeLayer = deepNet.AddReshapeLayer(depth, height, width);
    TReshapeLayer<Architecture_t> *copyReshapeLayer = new TReshapeLayer<Architecture_t>(*reshapeLayer);
 
    // add the copy to all slave nets