diff --git a/include/model/model.hpp b/include/model/model.hpp index 4a71e5f..fe50485 100644 --- a/include/model/model.hpp +++ b/include/model/model.hpp @@ -2,15 +2,24 @@ #define CUDANET_MODEL_H #include -#include #include +#include -#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; - std::unordered_map layerMap; + std::vector layers; + std::unordered_map layerMap; }; } // namespace CUDANet diff --git a/src/model/model.cpp b/src/model/model.cpp index cd02236..b3abc5e 100644 --- a/src/model/model.cpp +++ b/src/model/model.cpp @@ -1,5 +1,11 @@ #include "model.hpp" +#include +#include +#include +#include +#include + #include "input.cuh" #include "layer.cuh" @@ -25,7 +31,7 @@ Model::Model(const Model& other) outputLayer = new Layers::Output(*other.outputLayer); } -Model::~Model(){ +Model::~Model() { delete inputLayer; delete outputLayer; for (auto layer : layers) { @@ -51,4 +57,71 @@ void Model::addLayer(const std::string& name, Layers::SequentialLayer* layer) { if (wLayer != nullptr) { layerMap[name] = wLayer; } +} + +void Model::loadWeights(const std::string& path) { + std::ifstream file(path, std::ios::binary); + + if (!file.is_open()) { + std::cerr << "Failed to open file: " << path << std::endl; + return; + } + + int64_t headerSize; + file.read(reinterpret_cast(&headerSize), sizeof(headerSize)); + + std::string header(headerSize, '\0'); + file.read(&header[0], headerSize); + + std::vector tensorInfos; + size_t pos = 0; + + while (pos < header.size()) { + size_t nextPos = header.find('\n', pos); + if (nextPos == std::string::npos) + break; + + std::string line = header.substr(pos, nextPos - pos); + pos = nextPos + 1; + + size_t commaPos = line.find(','); + if (commaPos == std::string::npos) + continue; + + // Parse tensor name into name and type + std::string nameStr = line.substr(0, commaPos); + size_t dotPos = nameStr.find('.'); + if (dotPos == std::string::npos) + continue; + std::string name = nameStr.substr(0, dotPos); + TensorType type = nameStr.substr(dotPos + 1) == "w" ? TensorType::WEIGHT : TensorType::BIAS; + + line = line.substr(commaPos + 1); + + commaPos = line.find(','); + if (commaPos == std::string::npos) + continue; + + int size = std::stoi(line.substr(0, commaPos)); + int offset = std::stoi(line.substr(commaPos + 1)); + + tensorInfos.push_back({name, type, size, offset}); + } + + for (const auto& tensorInfo : tensorInfos) { + std::vector values(tensorInfo.size); + + file.seekg(tensorInfo.offset); + file.read(reinterpret_cast(values.data()), tensorInfo.size * sizeof(float)); + + if (layerMap.find(tensorInfo.name) != layerMap.end()) { + if (tensorInfo.type == TensorType::WEIGHT) { + layerMap[tensorInfo.name]->setWeights(values.data()); + } else if (tensorInfo.type == TensorType::BIAS) { + layerMap[tensorInfo.name]->setBiases(values.data()); + } + } + } + + file.close(); } \ No newline at end of file diff --git a/test/model/test_model.cu b/test/model/test_model.cu index 3d681cd..d306598 100644 --- a/test/model/test_model.cu +++ b/test/model/test_model.cu @@ -8,6 +8,8 @@ class ModelTest : public ::testing::Test { protected: CUDANet::Model *commonTestSetup( + bool setWeights = true, + int inputSize = 6, int inputChannels = 2, int outputSize = 3, @@ -28,17 +30,19 @@ class ModelTest : public ::testing::Test { CUDANet::Layers::Padding::VALID, CUDANet::Layers::ActivationType::NONE ); - // weights 6*6*2*2 - std::vector conv2dWeights = { - 0.18313f, 0.53363f, 0.39527f, 0.27575f, 0.3433f, 0.41746f, - 0.16831f, 0.61693f, 0.54599f, 0.99692f, 0.77127f, 0.25146f, - 0.4206f, 0.16291f, 0.93484f, 0.79765f, 0.74982f, 0.78336f, - 0.6386f, 0.87744f, 0.33587f, 0.9691f, 0.68437f, 0.65098f, - 0.48153f, 0.97546f, 0.8026f, 0.36689f, 0.98152f, 0.37351f, - 0.68407f, 0.2684f, 0.2855f, 0.76195f, 0.67828f, 0.603f - }; - conv2d->setWeights(conv2dWeights.data()); - model->addLayer("conv2d", conv2d); + + if (setWeights) { + std::vector conv2dWeights = { + 0.18313f, 0.53363f, 0.39527f, 0.27575f, 0.3433f, 0.41746f, + 0.16831f, 0.61693f, 0.54599f, 0.99692f, 0.77127f, 0.25146f, + 0.4206f, 0.16291f, 0.93484f, 0.79765f, 0.74982f, 0.78336f, + 0.6386f, 0.87744f, 0.33587f, 0.9691f, 0.68437f, 0.65098f, + 0.48153f, 0.97546f, 0.8026f, 0.36689f, 0.98152f, 0.37351f, + 0.68407f, 0.2684f, 0.2855f, 0.76195f, 0.67828f, 0.603f + }; + conv2d->setWeights(conv2dWeights.data()); + } + model->addLayer("conv1", conv2d); // maxpool2d CUDANet::Layers::MaxPooling2D *maxpool2d = @@ -46,27 +50,29 @@ class ModelTest : public ::testing::Test { inputSize - kernelSize + 1, numFilters, poolingSize, poolingStride, CUDANet::Layers::ActivationType::RELU ); - model->addLayer("maxpool2d", maxpool2d); + model->addLayer("maxpool1", maxpool2d); // dense CUDANet::Layers::Dense *dense = new CUDANet::Layers::Dense( 8, 3, CUDANet::Layers::ActivationType::SOFTMAX ); - // dense weights 18*6 - std::vector denseWeights = { - 0.36032f, 0.33115f, 0.02948f, 0.09802f, 0.45072f, 0.56266f, - 0.43514f, 0.80946f, 0.43439f, 0.90916f, 0.08605f, 0.07473f, - 0.94788f, 0.66168f, 0.34927f, 0.09464f, 0.61963f, 0.73775f, - 0.51559f, 0.81916f, 0.64915f, 0.03934f, 0.87608f, 0.68364f, - }; - dense->setWeights(denseWeights.data()); - model->addLayer("dense", dense); + + if (setWeights) { + std::vector denseWeights = { + 0.36032f, 0.33115f, 0.02948f, 0.09802f, 0.45072f, 0.56266f, + 0.43514f, 0.80946f, 0.43439f, 0.90916f, 0.08605f, 0.07473f, + 0.94788f, 0.66168f, 0.34927f, 0.09464f, 0.61963f, 0.73775f, + 0.51559f, 0.81916f, 0.64915f, 0.03934f, 0.87608f, 0.68364f, + }; + dense->setWeights(denseWeights.data()); + } + model->addLayer("linear", dense); return model; } void commonTestTeardown(CUDANet::Model *model) { - delete model; + delete model; } cudaError_t cudaStatus; @@ -103,6 +109,8 @@ TEST_F(ModelTest, TestModelPredict) { } EXPECT_NEAR(sum, 1.0f, 1e-5f); + + commonTestTeardown(model); } TEST_F(ModelTest, TestModelPredictMultiple) { @@ -162,4 +170,43 @@ TEST_F(ModelTest, TestModelPredictMultiple) { } EXPECT_NEAR(sum_2, 1.0f, 1e-5f); + + commonTestTeardown(model); +} + +TEST_F(ModelTest, TestLoadWeights) { + int outputSize = 3; + CUDANet::Model *model = commonTestSetup(); + + model->loadWeights("../test/resources/model.bin"); + + std::vector input = { + 0.12762f, 0.99056f, 0.77565f, 0.29058f, 0.29787f, 0.58415f, 0.20484f, + 0.05415f, 0.60593f, 0.3162f, 0.08198f, 0.92749f, 0.72392f, 0.91786f, + 0.65266f, 0.80908f, 0.53389f, 0.36069f, 0.18614f, 0.52381f, 0.08525f, + 0.43054f, 0.3355f, 0.96587f, 0.98194f, 0.71336f, 0.78392f, 0.50648f, + 0.40355f, 0.31863f, 0.54686f, 0.1836f, 0.77171f, 0.01262f, 0.41108f, + 0.53467f, 0.3553f, 0.42808f, 0.45798f, 0.29958f, 0.3923f, 0.98277f, + 0.02033f, 0.99868f, 0.90584f, 0.57554f, 0.15957f, 0.91273f, 0.38901f, + 0.27097f, 0.64788f, 0.84272f, 0.42984f, 0.07466f, 0.53658f, 0.83388f, + 0.28232f, 0.48046f, 0.85626f, 0.04721f, 0.36139f, 0.6123f, 0.56991f, + 0.84854f, 0.61415f, 0.2466f, 0.20017f, 0.78952f, 0.93797f, 0.27884f, + 0.30514f, 0.23521f + }; + + std::vector expected = {2e-05f, 0.00021f, 0.99977f}; + + // predict + const float *output = model->predict(input.data()); + + float sum = 0.0f; + for (int i = 0; i < outputSize; ++i) { + EXPECT_NEAR(expected[i], output[i], 1e-5f); + sum += output[i]; + } + + EXPECT_NEAR(sum, 1.0f, 1e-5f); + + commonTestTeardown(model); + } \ No newline at end of file