diff --git a/README.md b/README.md index ddb35b9..a3c7907 100644 --- a/README.md +++ b/README.md @@ -4,20 +4,7 @@ Convolutional Neural Network inference library running on CUDA. -## Features - -- [x] Input layer -- [x] Dense (fully-connected) layer -- [x] Conv2d layer -- [x] Max pooling -- [x] Average pooling -- [x] Concat layer -- [x] Sigmoid activation -- [x] ReLU activation -- [x] Softmax activation -- [x] Load weights from file - -## Usage +## Quickstart Guide **requirements** - [cmake](https://cmake.org/) @@ -94,7 +81,7 @@ float* MyModel::predict(const float* input) { CUDANet uses format similar to safetensors to load weights and biases. ``` -[int64 header size, header, tensor values] +[u_short version, u_int64 header size, header, tensor values] ``` where `header` is a csv format @@ -103,6 +90,4 @@ where `header` is a csv format ,, ``` -To load weights call `load_weights` function on Model object. - -To export weights from pytorch you can use the `export_model_weights` function from `tools/utils.py` script \ No newline at end of file +To load weights call `load_weights` function on Model object. To export weights from pytorch you can use the `export_model_weights` function from `tools/utils.py` script. Currently only float32 weights are supported \ No newline at end of file diff --git a/src/model/model.cpp b/src/model/model.cpp index a8530f6..e992f14 100644 --- a/src/model/model.cpp +++ b/src/model/model.cpp @@ -66,7 +66,15 @@ void Model::loadWeights(const std::string& path) { return; } - int64_t headerSize; + u_short version; + file.read(reinterpret_cast(&version), sizeof(version)); + + if (version != 1) { + std::cerr << "Unsupported model version: " << version << std::endl; + return; + } + + u_int64_t headerSize; file.read(reinterpret_cast(&headerSize), sizeof(headerSize)); std::string header(headerSize, '\0'); @@ -110,7 +118,7 @@ void Model::loadWeights(const std::string& path) { for (const auto& tensorInfo : tensorInfos) { std::vector values(tensorInfo.size); - file.seekg(sizeof(int64_t) + header.size() + tensorInfo.offset); + file.seekg(sizeof(version) + sizeof(headerSize) + header.size() + tensorInfo.offset); file.read(reinterpret_cast(values.data()), tensorInfo.size * sizeof(float)); if (layerMap.find(tensorInfo.name) != layerMap.end()) { diff --git a/test/resources/model.bin b/test/resources/model.bin index d58f27a..2f2ff22 100644 Binary files a/test/resources/model.bin and b/test/resources/model.bin differ diff --git a/tools/utils.py b/tools/utils.py index 6c4f689..95da2b8 100644 --- a/tools/utils.py +++ b/tools/utils.py @@ -16,6 +16,7 @@ def print_cpp_vector(vector, name="expected"): def export_model_weights(model: torch.nn.Module, filename): with open(filename, 'wb') as f: + version = 1 header = "" offset = 0 tensor_data = b"" @@ -33,7 +34,8 @@ def export_model_weights(model: torch.nn.Module, filename): tensor_data += tensor_bytes f.seek(0) - f.write(struct.pack('q', len(header))) + f.write(struct.pack('H', version)) + f.write(struct.pack('Q', len(header))) f.write(header.encode('utf-8')) f.write(tensor_data)