From 942ee6a32b23eb0360a87bef8e588f333cc5f61d Mon Sep 17 00:00:00 2001
From: LordMathis <matus@namesny.com>
Date: Sun, 21 Apr 2024 12:20:02 +0200
Subject: [PATCH] Add layer name to vector

---
 include/model/model.hpp |  9 ++++++---
 src/model/model.cpp     | 14 +++++++-------
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/include/model/model.hpp b/include/model/model.hpp
index 98fc392..126c7a5 100644
--- a/include/model/model.hpp
+++ b/include/model/model.hpp
@@ -11,7 +11,10 @@
 
 namespace CUDANet {
 
-enum TensorType { WEIGHT, BIAS, };
+enum TensorType {
+    WEIGHT,
+    BIAS,
+};
 
 struct TensorInfo {
     std::string name;
@@ -42,8 +45,8 @@ class Model {
 
     int outputSize;
 
-    std::vector<Layers::SequentialLayer*>                   layers;
-    std::unordered_map<std::string, Layers::SequentialLayer*> layerMap;
+    std::vector<std::pair<std::string, Layers::SequentialLayer*>> layers;
+    std::unordered_map<std::string, Layers::SequentialLayer*>     layerMap;
 };
 
 }  // namespace CUDANet
diff --git a/src/model/model.cpp b/src/model/model.cpp
index 07c721d..408c1e9 100644
--- a/src/model/model.cpp
+++ b/src/model/model.cpp
@@ -15,7 +15,7 @@ Model::Model(const int inputSize, const int inputChannels, const int outputSize)
     : inputSize(inputSize),
       inputChannels(inputChannels),
       outputSize(outputSize),
-      layers(std::vector<Layers::SequentialLayer*>()),
+      layers(std::vector<std::pair<std::string, Layers::SequentialLayer*>>()),
       layerMap(std::unordered_map<std::string, Layers::SequentialLayer*>()) {
     inputLayer  = new Layers::Input(inputSize * inputSize * inputChannels);
     outputLayer = new Layers::Output(outputSize);
@@ -25,7 +25,7 @@ Model::Model(const Model& other)
     : inputSize(other.inputSize),
       inputChannels(other.inputChannels),
       outputSize(other.outputSize),
-      layers(std::vector<Layers::SequentialLayer*>()),
+      layers(std::vector<std::pair<std::string, Layers::SequentialLayer*>>()),
       layerMap(std::unordered_map<std::string, Layers::SequentialLayer*>()) {
     inputLayer  = new Layers::Input(*other.inputLayer);
     outputLayer = new Layers::Output(*other.outputLayer);
@@ -34,8 +34,8 @@ Model::Model(const Model& other)
 Model::~Model() {
     delete inputLayer;
     delete outputLayer;
-    for (auto layer : layers) {
-        delete layer;
+    for (const auto& layer : layers) {
+        delete layer.second;
     }
 };
 
@@ -43,15 +43,15 @@ float* Model::predict(const float* input) {
     float* d_input = inputLayer->forward(input);
 
     for (auto& layer : layers) {
-        std::cout << layer << std::endl;
-        d_input = layer->forward(d_input);
+        std::cout << layer.first << std::endl;
+        d_input = layer.second->forward(d_input);
     }
 
     return outputLayer->forward(d_input);
 }
 
 void Model::addLayer(const std::string& name, Layers::SequentialLayer* layer) {
-    layers.push_back(layer);
+    layers.push_back({ name, layer });
     layerMap[name] = layer;
 }