Implement modules

include/model/model.hpp

@@ -7,6 +7,7 @@
 
 #include "input.cuh"
 #include "layer.cuh"
+#include "module.hpp"
 #include "output.cuh"
 
 namespace CUDANet {

include/model/module.hpp

@@ -0,0 +1,42 @@
+#ifndef CUDANET_MODULE_H
+#define CUDANET_MODULE_H
+
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "layer.cuh"
+
+namespace CUDANet {
+
+class Module : public Layers::SequentialLayer {
+  public:
+    Module(const int inputSize, const int inputChannels, const int outputSize, const int outputChannels);
+    ~Module();
+
+    virtual float* forward(const float* d_input) = 0;
+
+    int getOutputSize();
+    int getInputSize();
+
+    void addLayer(const std::string& name, Layers::SequentialLayer* layer);
+    Layers::SequentialLayer* getLayer(const std::string& name);
+
+    const std::unordered_map<std::string, Layers::SequentialLayer*>& getLayers() const;
+
+  private:
+    int inputSize;
+    int inputChannels;
+
+    int outputSize;
+    int outputChannels;
+
+    float *d_output;
+
+    std::vector<std::pair<std::string, Layers::SequentialLayer*>> layers;
+    std::unordered_map<std::string, Layers::SequentialLayer*>     layerMap;
+};
+
+}  // namespace CUDANet
+
+#endif

src/model/model.cpp

@@ -50,6 +50,18 @@ float* Model::predict(const float* input) {
 }
 
 void Model::addLayer(const std::string& name, Layers::SequentialLayer* layer) {
+
+    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;
+        }
+
+        return;
+    }
+
     layers.push_back({ name, layer });
     layerMap[name] = layer;
 }

src/model/module.cpp

@@ -0,0 +1,49 @@
+#include "module.hpp"
+
+#include "cuda_helper.cuh"
+
+using namespace CUDANet;
+
+Module::Module(
+    const int inputSize,
+    const int inputChannels,
+    const int outputSize,
+    const int outputChannels
+)
+    : inputSize(inputSize),
+      inputChannels(inputChannels),
+      outputSize(outputSize),
+      outputChannels(outputChannels),
+      layers(std::vector<std::pair<std::string, Layers::SequentialLayer*>>()),
+      layerMap(std::unordered_map<std::string, Layers::SequentialLayer*>()) {
+    d_output = nullptr;
+    CUDA_CHECK(cudaMalloc(
+        (void**)&d_output,
+        sizeof(float) * outputSize * outputSize * outputChannels
+    ));
+}
+
+Module::~Module() {
+    cudaFree(d_output);
+}
+
+void Module::addLayer(const std::string& name, Layers::SequentialLayer* layer) {
+    layers.push_back({ name, layer });
+    layerMap[name] = layer;
+}
+
+Layers::SequentialLayer* Module::getLayer(const std::string& name) {
+    return layerMap[name];
+}
+
+const std::unordered_map<std::string, Layers::SequentialLayer*>& Module::getLayers() const {
+    return layerMap;
+}
+
+int Module::getInputSize() {
+    return inputSize * inputSize * inputChannels;
+}
+
+int Module::getOutputSize() {
+    return outputSize * outputSize * outputChannels;
+}