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

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

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