Add output layer to model predict

This commit is contained in:
2024-03-21 23:22:12 +01:00
parent 90fb104dae
commit 87db47089e
2 changed files with 32 additions and 18 deletions

View File

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