From 87db47089e9c7b37fee7af3ddd1fed366c906c8c Mon Sep 17 00:00:00 2001
From: LordMathis <matus@namesny.com>
Date: Thu, 21 Mar 2024 23:22:12 +0100
Subject: [PATCH] Add output layer to model predict

---
 include/model/model.hpp |  9 ++++++---
 src/model/model.cpp     | 41 ++++++++++++++++++++++++++---------------
 2 files changed, 32 insertions(+), 18 deletions(-)

diff --git a/include/model/model.hpp b/include/model/model.hpp
index be7711e..4a71e5f 100644
--- a/include/model/model.hpp
+++ b/include/model/model.hpp
@@ -3,16 +3,18 @@
 
 #include <string>
 #include <vector>
-#include <map>
+#include <unordered_map>
 
 #include "layer.cuh"
 #include "input.cuh"
+#include "output.cuh"
 
 namespace CUDANet {
 
 class Model {
   public:
-    Model(const int inputSize, const int inputChannels);
+    Model(const int inputSize, const int inputChannels, const int outputSize);
+    Model(const Model& other);
     ~Model();
 
     float* predict(const float* input);
@@ -22,6 +24,7 @@ class Model {
   private:
 
     Layers::Input *inputLayer;
+    Layers::Output *outputLayer;
 
     int inputSize;
     int inputChannels;
@@ -29,7 +32,7 @@ class Model {
     int outputSize;
     
     std::vector<Layers::SequentialLayer*> layers;
-    std::map<std::string, Layers::WeightedLayer*> layerMap;
+    std::unordered_map<std::string, Layers::WeightedLayer*> layerMap;
 
 };
 
diff --git a/src/model/model.cpp b/src/model/model.cpp
index 82f441f..46fd799 100644
--- a/src/model/model.cpp
+++ b/src/model/model.cpp
@@ -1,37 +1,48 @@
 #include "model.hpp"
-#include "layer.cuh"
+
 #include "input.cuh"
+#include "layer.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;
-    inputLayer = new Layers::Input(inputLayerSize);
+Model::Model(const int inputSize, const int inputChannels, const int outputSize)
+    : inputSize(inputSize),
+      inputChannels(inputChannels),
+      outputSize(outputSize),
+      layers(std::vector<Layers::SequentialLayer*>()),
+      layerMap(std::unordered_map<std::string, Layers::WeightedLayer*>()) {
+    inputLayer  = new Layers::Input(inputSize * inputSize * inputChannels);
+    outputLayer = new Layers::Output(outputSize);
 };
 
+Model::Model(const Model& other)
+    : inputSize(other.inputSize),
+      inputChannels(other.inputChannels),
+      outputSize(other.outputSize),
+      layers(std::vector<Layers::SequentialLayer*>()),
+      layerMap(std::unordered_map<std::string, Layers::WeightedLayer*>()) {
+    inputLayer  = new Layers::Input(*other.inputLayer);
+    outputLayer = new Layers::Output(*other.outputLayer);
+}
+
 Model::~Model(){};
 
 float* Model::predict(const float* input) {
-
     float* d_input = inputLayer->forward(input);
-    
+
     for (auto& layer : layers) {
         d_input = layer->forward(d_input);
     }
 
-    return d_input;
+    return outputLayer->forward(d_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);
+    Layers::WeightedLayer* wLayer = dynamic_cast<Layers::WeightedLayer*>(layer);
+
+    if (wLayer != nullptr) {
+        layerMap[name] = wLayer;
     }
 }
\ No newline at end of file