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 <string>
#include <vector> #include <vector>
#include <map> #include <unordered_map>
#include "layer.cuh" #include "layer.cuh"
#include "input.cuh" #include "input.cuh"
#include "output.cuh"
namespace CUDANet { namespace CUDANet {
class Model { class Model {
public: public:
Model(const int inputSize, const int inputChannels); Model(const int inputSize, const int inputChannels, const int outputSize);
Model(const Model& other);
~Model(); ~Model();
float* predict(const float* input); float* predict(const float* input);
@@ -22,6 +24,7 @@ class Model {
private: private:
Layers::Input *inputLayer; Layers::Input *inputLayer;
Layers::Output *outputLayer;
int inputSize; int inputSize;
int inputChannels; int inputChannels;
@@ -29,7 +32,7 @@ class Model {
int outputSize; int outputSize;
std::vector<Layers::SequentialLayer*> layers; 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 "model.hpp"
#include "layer.cuh"
#include "input.cuh" #include "input.cuh"
#include "layer.cuh"
using namespace CUDANet; using namespace CUDANet;
Model::Model(const int inputSize, const int inputChannels) Model::Model(const int inputSize, const int inputChannels, const int outputSize)
: inputSize(inputSize), inputChannels(inputChannels) { : inputSize(inputSize),
inputChannels(inputChannels),
layerMap = std::map<std::string, Layers::WeightedLayer*>(); outputSize(outputSize),
layers = std::vector<Layers::SequentialLayer*>(); layers(std::vector<Layers::SequentialLayer*>()),
layerMap(std::unordered_map<std::string, Layers::WeightedLayer*>()) {
inputLayer = new Layers::Input(inputSize * inputSize * inputChannels);
const int inputLayerSize = inputSize * inputSize * inputChannels; outputLayer = new Layers::Output(outputSize);
inputLayer = new Layers::Input(inputLayerSize);
}; };
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(){}; Model::~Model(){};
float* Model::predict(const float* input) { float* Model::predict(const float* input) {
float* d_input = inputLayer->forward(input); float* d_input = inputLayer->forward(input);
for (auto& layer : layers) { for (auto& layer : layers) {
d_input = layer->forward(d_input); d_input = layer->forward(d_input);
} }
return d_input; return outputLayer->forward(d_input);
} }
void Model::addLayer(const std::string& name, Layers::SequentialLayer* layer) { void Model::addLayer(const std::string& name, Layers::SequentialLayer* layer) {
layers.push_back(layer); layers.push_back(layer);
if (dynamic_cast<Layers::WeightedLayer*>(layer) != nullptr) { Layers::WeightedLayer* wLayer = dynamic_cast<Layers::WeightedLayer*>(layer);
layerMap[name] = dynamic_cast<Layers::WeightedLayer*>(layer);
if (wLayer != nullptr) {
layerMap[name] = wLayer;
} }
} }