From 8ac2da004c8c3768c06344877f1ac58f94abeffa Mon Sep 17 00:00:00 2001
From: LordMathis <matus@namesny.com>
Date: Wed, 29 May 2024 21:02:26 +0200
Subject: [PATCH] Fix adding module layers

---
 examples/inception_v3/inception_v3.cpp | 32 +++++++++++++++-----------
 include/model/module.hpp               |  2 --
 src/model/model.cpp                    |  5 ++--
 src/model/module.cpp                   | 12 +++++++---
 4 files changed, 31 insertions(+), 20 deletions(-)

diff --git a/examples/inception_v3/inception_v3.cpp b/examples/inception_v3/inception_v3.cpp
index 9ec68b1..f74ea9c 100644
--- a/examples/inception_v3/inception_v3.cpp
+++ b/examples/inception_v3/inception_v3.cpp
@@ -1,16 +1,6 @@
 #include <cudanet.cuh>
 #include <iostream>
 
-int main(int argc, const char *const argv[]) {
-    if (argc != 3) {
-        std::cerr << "Usage: " << argv[0] << "<model_weights_path> <image_path>"
-                  << std::endl;
-        return 1;  // Return error code indicating incorrect usage
-    }
-
-    std::cout << "Loading model..." << std::endl;
-}
-
 class BasicConv2d : public CUDANet::Module {
   public:
     BasicConv2d(
@@ -702,7 +692,7 @@ class InceptionE : public CUDANet::Module {
     }
 
     shape2d getOutputDims() {
-        branch3x3_2a->getOutputDims();
+        return branch3x3_2a->getOutputDims();
     }
 
     int getOutputChannels() {
@@ -823,7 +813,7 @@ class InceptionV3 : public CUDANet::Model {
             Mixed_7c->getOutputSize(), 1000,
             CUDANet::Layers::ActivationType::SOFTMAX
         );
-        addLayer("", fc);
+        addLayer("fc", fc);
     }
 
     float* predict(const float* input) {
@@ -905,4 +895,20 @@ class InceptionV3 : public CUDANet::Model {
     InceptionE *Mixed_7c;
 
     CUDANet::Layers::Dense *fc;
-};
\ No newline at end of file
+};
+
+int main(int argc, const char *const argv[]) {
+
+    InceptionV3 *inception_v3 = new InceptionV3({299, 299}, 3, 1000);
+
+    inception_v3->printSummary();
+
+
+    if (argc != 3) {
+        std::cerr << "Usage: " << argv[0] << "<model_weights_path> <image_path>"
+                  << std::endl;
+        return 1;  // Return error code indicating incorrect usage
+    }
+
+    std::cout << "Loading model..." << std::endl;
+}
\ No newline at end of file
diff --git a/include/model/module.hpp b/include/model/module.hpp
index 6fd5ada..4609b20 100644
--- a/include/model/module.hpp
+++ b/include/model/module.hpp
@@ -11,8 +11,6 @@ namespace CUDANet {
 
 class Module : public Layers::SequentialLayer {
   public:
-    virtual ~Module() = 0;
-
     virtual float* forward(const float* d_input) = 0;
 
     int getOutputSize();
diff --git a/src/model/model.cpp b/src/model/model.cpp
index efd880c..5386edd 100644
--- a/src/model/model.cpp
+++ b/src/model/model.cpp
@@ -56,17 +56,18 @@ float* Model::predict(const float* input) {
 }
 
 void Model::addLayer(const std::string& name, Layers::SequentialLayer* layer) {
-    Module* module = dynamic_cast<Module*>(layer);
+    const Module* module = dynamic_cast<Module*>(layer);
 
     if (module != nullptr) {
-        layers.push_back({name, module});
         for (const auto& moduleLayer : module->getLayers()) {
             layerMap[moduleLayer.first] = moduleLayer.second;
+            layers.push_back({moduleLayer.first, moduleLayer.second});
         }
 
         return;
     }
 
+
     layers.push_back({name, layer});
     layerMap[name] = layer;
 }
diff --git a/src/model/module.cpp b/src/model/module.cpp
index 4fd2728..ea82e6c 100644
--- a/src/model/module.cpp
+++ b/src/model/module.cpp
@@ -5,23 +5,29 @@
 using namespace CUDANet;
 
 void Module::addLayer(const std::string& name, Layers::SequentialLayer* layer) {
-    Module* module = dynamic_cast<Module*>(layer);
+    const Module* module = dynamic_cast<Module*>(layer);
 
     if (module != nullptr) {
-        layers.push_back({ name, module });
         for (const auto& moduleLayer : module->getLayers()) {
             layerMap[moduleLayer.first] = moduleLayer.second;
+            layers.push_back({moduleLayer.first, moduleLayer.second});
         }
 
         return;
     }
+
+    layers.push_back({name, layer});
+    layerMap[name] = layer;
+
+    // std::cout << "Wat?! - module" << name << std::endl;
 }
 
 Layers::SequentialLayer* Module::getLayer(const std::string& name) {
     return layerMap[name];
 }
 
-const std::unordered_map<std::string, Layers::SequentialLayer*>& Module::getLayers() const {
+const std::unordered_map<std::string, Layers::SequentialLayer*>&
+Module::getLayers() const {
     return layerMap;
 }