fix: adam not considering Flatten layer

Az-r-ow · Jul 8, 2024 · c05c9e9 · c05c9e9
1 parent 5927c91
commit c05c9e9
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
diff --git a/src/NeuralNet/Network.cpp b/src/NeuralNet/Network.cpp
@@ -351,11 +351,12 @@ void Network::updateOptimizerSetup(size_t numLayers) {
 
   if (std::dynamic_pointer_cast<Adam>(this->optimizer)) {
     // Get the number of dense layers
-    nLayers =
-        std::count_if(this->layers.begin(), this->layers.end(),
-                      [](const std::shared_ptr<Layer> &ptr) {
-                        return std::dynamic_pointer_cast<Dense>(ptr) != nullptr;
-                      });
+    nLayers = std::count_if(
+        this->layers.begin(), this->layers.end(),
+        [](const std::shared_ptr<Layer> &ptr) {
+          return std::dynamic_pointer_cast<Dense>(ptr) != nullptr ||
+                 std::dynamic_pointer_cast<Flatten>(ptr) != nullptr;
+        });
   }
 
   this->optimizer->insiderInit(nLayers);

diff --git a/src/NeuralNet/optimizers/Adam.hpp b/src/NeuralNet/optimizers/Adam.hpp
@@ -58,6 +58,9 @@ class Adam : public Optimizer {
       v = Eigen::MatrixBase<Derived1>::Zero(param.rows(), param.cols());
     }
 
+    assert(gradients.rows() == m.rows() && gradients.cols() == m.cols());
+    assert(gradients.rows() == v.rows() && gradients.cols() == v.cols());
+
     // update biased first moment estimate
     m = (beta1 * m).array() + ((1 - beta2) * gradients.array()).array();