Mathis/CUDANet
1
0
Fork 0
mirror of https://github.com/lordmathis/CUDANet.git synced 2025-11-05 17:34:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add output layer to model predict

Browse Source
This commit is contained in:
Matúš Námešný
2024-03-21 23:22:12 +01:00
parent 90fb104dae
commit 87db47089e
2 changed files with 32 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
include/model/model.hpp
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;
};