From af6838e8ae5d5cebc43ac5377b9c352b9440b8ef Mon Sep 17 00:00:00 2001
From: LordMathis <matus@namesny.com>
Date: Wed, 20 Mar 2024 22:31:39 +0100
Subject: [PATCH] Initial model implementation

---
 CMakeLists.txt           |  2 ++
 include/layers/input.cuh |  6 ++++--
 include/model/model.hpp  | 34 ++++++++++++++++++++++++++++++++++
 src/model/model.cpp      | 36 ++++++++++++++++++++++++++++++++++++
 4 files changed, 76 insertions(+), 2 deletions(-)
 create mode 100644 include/model/model.hpp
 create mode 100644 src/model/model.cpp

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 792d825..bfcd083 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -12,6 +12,7 @@ file(GLOB_RECURSE LIBRARY_SOURCES
     src/utils/*.cu
     src/kernels/*.cu
     src/layers/*.cu
+    src/model/*.cpp
 )
 
 set(LIBRARY_SOURCES
@@ -32,6 +33,7 @@ target_include_directories(${PROJECT_NAME} PUBLIC
     ${CMAKE_CURRENT_SOURCE_DIR}/include/utils
     ${CMAKE_CURRENT_SOURCE_DIR}/include/kernels
     ${CMAKE_CURRENT_SOURCE_DIR}/include/layers
+    ${CMAKE_CURRENT_SOURCE_DIR}/include/model
     ${CMAKE_CURRENT_SOURCE_DIR}/src
 )
 
diff --git a/include/layers/input.cuh b/include/layers/input.cuh
index 33fdb1d..144cbe9 100644
--- a/include/layers/input.cuh
+++ b/include/layers/input.cuh
@@ -1,20 +1,22 @@
 #ifndef CUDANET_INPUT_LAYER_H
 #define CUDANET_INPUT_LAYER_H
 
+#include "layer.cuh"
+
 namespace CUDANet::Layers {
 
 /**
  * @brief Input layer, just copies the input to the device
  *
  */
-class Input {
+class Input : public SequentialLayer {
   public:
     /**
      * @brief Create a new Input layer
      * 
      * @param inputSize Size of the input vector
      */
-    Input(int inputSize);
+    explicit Input(int inputSize);
 
     /**
      * @brief Destroy the Input layer
diff --git a/include/model/model.hpp b/include/model/model.hpp
new file mode 100644
index 0000000..406cb8e
--- /dev/null
+++ b/include/model/model.hpp
@@ -0,0 +1,34 @@
+#ifndef CUDANET_MODEL_H
+#define CUDANET_MODEL_H
+
+#include <string>
+#include <vector>
+#include <map>
+#include "layer.cuh"
+
+namespace CUDANet {
+
+class Model {
+  public:
+    Model(const int inputSize, const int inputChannels);
+    ~Model();
+
+    float* predict(const float* input);
+
+    void addLayer(const std::string& name, Layers::SequentialLayer* layer);
+
+  private:
+
+    int inputSize;
+    int inputChannels;
+
+    int outputSize;
+    
+    std::vector<Layers::SequentialLayer*> layers;
+    std::map<std::string, Layers::WeightedLayer*> layerMap;
+
+};
+
+}  // namespace CUDANet
+
+#endif  // CUDANET_MODEL_H
\ No newline at end of file
diff --git a/src/model/model.cpp b/src/model/model.cpp
new file mode 100644
index 0000000..0291e8b
--- /dev/null
+++ b/src/model/model.cpp
@@ -0,0 +1,36 @@
+#include "model.hpp"
+#include "layer.cuh"
+#include "input.cuh"
+
+using namespace CUDANet;
+
+Model::Model(const int inputSize, const int inputChannels)
+    : inputSize(inputSize), inputChannels(inputChannels) {
+
+    layerMap = std::map<std::string, Layers::WeightedLayer*>();
+    layers   = std::vector<Layers::SequentialLayer*>();
+
+
+    const int inputLayerSize = inputSize * inputSize * inputChannels;
+    Layers::Input* inputLayer = new Layers::Input(inputLayerSize);
+    
+    layers.push_back(inputLayer);
+};
+
+Model::~Model(){};
+
+float* Model::predict(const float* input) {
+    
+    for (auto& layer : layers) {
+        input = layer->forward(input);
+    }
+
+}
+
+void Model::addLayer(const std::string& name, Layers::SequentialLayer* layer) {
+    layers.push_back(layer);
+
+    if (dynamic_cast<Layers::WeightedLayer*>(layer) != nullptr) {
+        layerMap[name] = dynamic_cast<Layers::WeightedLayer*>(layer);
+    }
+}
\ No newline at end of file