diff --git a/examples/alexnet/main.cpp b/examples/alexnet/main.cpp index 1e88b9d..e133f65 100644 --- a/examples/alexnet/main.cpp +++ b/examples/alexnet/main.cpp @@ -109,6 +109,9 @@ int main(int argc, const char* const argv[]) { const int outputSize = 1000; CUDANet::Model *model = createModel(inputSize, inputChannels, outputSize); + + model->validate(); + model->loadWeights(modelWeightsPath); // Read and normalize the image diff --git a/include/model/model.hpp b/include/model/model.hpp index 126c7a5..cdb1158 100644 --- a/include/model/model.hpp +++ b/include/model/model.hpp @@ -36,6 +36,8 @@ class Model { void loadWeights(const std::string& path); + bool validate(); + private: Layers::Input* inputLayer; Layers::Output* outputLayer; diff --git a/src/model/model.cpp b/src/model/model.cpp index 8b138c9..a8530f6 100644 --- a/src/model/model.cpp +++ b/src/model/model.cpp @@ -147,4 +147,23 @@ void Model::loadWeights(const std::string& path) { } file.close(); +} + +bool Model::validate() { + + bool valid = true; + int size = inputLayer->getInputSize(); + + for (const auto& layer : layers) { + if (layer.second->getInputSize() != size) { + valid = false; + std::cerr << "Layer: " << layer.first << " has incorrect input size, expected " << size << " but got " + << layer.second->getInputSize() << std::endl; + break; + } + + size = layer.second->getOutputSize(); + } + + return valid; } \ No newline at end of file