Mathis/CUDANet
1
0
Fork 0
mirror of https://github.com/lordmathis/CUDANet.git synced 2025-11-06 09:44:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix adding module layers

Browse Source
This commit is contained in:
Matúš Námešný
2024-05-29 21:02:26 +02:00
parent 4159454842
commit 8ac2da004c
4 changed files with 31 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
src/model/model.cpp
View File

@@ -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;
}

12
src/model/module.cpp
View File

@@ -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;
}