Start implementing weights import

This commit is contained in:
2024-04-15 22:17:48 +02:00
parent d8c50116e8
commit f4ae45f867
3 changed files with 159 additions and 31 deletions

View File

@@ -2,15 +2,24 @@
#define CUDANET_MODEL_H
#include <string>
#include <vector>
#include <unordered_map>
#include <vector>
#include "layer.cuh"
#include "input.cuh"
#include "layer.cuh"
#include "output.cuh"
namespace CUDANet {
enum TensorType { WEIGHT, BIAS, };
struct TensorInfo {
std::string name;
TensorType type;
int size;
int offset;
};
class Model {
public:
Model(const int inputSize, const int inputChannels, const int outputSize);
@@ -20,20 +29,19 @@ class Model {
float* predict(const float* input);
void addLayer(const std::string& name, Layers::SequentialLayer* layer);
void loadWeights(const std::string& path);
private:
Layers::Input *inputLayer;
Layers::Output *outputLayer;
Layers::Input* inputLayer;
Layers::Output* outputLayer;
int inputSize;
int inputChannels;
int outputSize;
std::vector<Layers::SequentialLayer*> layers;
std::unordered_map<std::string, Layers::WeightedLayer*> layerMap;
std::vector<Layers::SequentialLayer*> layers;
std::unordered_map<std::string, Layers::WeightedLayer*> layerMap;
};
} // namespace CUDANet