From 0807a0f2b8d88049135bb2b75a63a7fd2f4a4434 Mon Sep 17 00:00:00 2001
From: LordMathis <matus@namesny.com>
Date: Sat, 20 Apr 2024 15:49:59 +0200
Subject: [PATCH] Export pretrained alexnet

---
 README.md                     |  2 +-
 examples/alexnet/alexnet.py   | 13 ++++++++++++
 examples/alexnet/main.cpp     | 12 +++++++++++
 tools/export_model_weights.py | 28 -------------------------
 tools/utils.py                | 39 +++++++++++++++++++++++++++++++++--
 5 files changed, 63 insertions(+), 31 deletions(-)
 create mode 100644 examples/alexnet/alexnet.py
 delete mode 100644 tools/export_model_weights.py

diff --git a/README.md b/README.md
index a3d178e..ddb35b9 100644
--- a/README.md
+++ b/README.md
@@ -105,4 +105,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 `tools/export_model_weights.py` script
\ No newline at end of file
+To export weights from pytorch you can use the `export_model_weights` function from `tools/utils.py`  script
\ No newline at end of file
diff --git a/examples/alexnet/alexnet.py b/examples/alexnet/alexnet.py
new file mode 100644
index 0000000..d1fcd56
--- /dev/null
+++ b/examples/alexnet/alexnet.py
@@ -0,0 +1,13 @@
+import torchvision
+import sys
+
+sys.path.append('../../tools')  # Ugly hack 
+from utils import export_model_weights, print_model_parameters
+
+if __name__ == "__main__":
+    alexnet = torchvision.models.alexnet(pretrained=True)
+    print_model_parameters(alexnet)  # print layer names and number of parameters
+    export_model_weights(alexnet, 'alexnet_weights.bin')
+    print()
+    print(alexnet)
+
diff --git a/examples/alexnet/main.cpp b/examples/alexnet/main.cpp
index a15cb46..05fc679 100644
--- a/examples/alexnet/main.cpp
+++ b/examples/alexnet/main.cpp
@@ -4,6 +4,7 @@
 #include <opencv2/opencv.hpp>
 
 #include <model.hpp>
+#include <conv2d.cuh>
 
 std::vector<float> readAndNormalizeImage(const std::string& imagePath, int width, int height) {
     // Read the image using OpenCV
@@ -28,6 +29,17 @@ std::vector<float> readAndNormalizeImage(const std::string& imagePath, int width
 CUDANet::Model* createModel(const int inputSize, const int inputChannels, const int outputSize) {
     CUDANet::Model *model =
         new CUDANet::Model(inputSize, inputChannels, outputSize);
+
+    // AlexNet
+    CUDANet::Layers::Conv2d *conv1 = new CUDANet::Layers::Conv2d(
+        inputSize, inputChannels, 11, 4, 96, CUDANet::Layers::Padding::SAME, CUDANet::Layers::ActivationType::RELU
+    );
+    model->addLayer("conv1", conv1);
+    CUDANet::Layers::MaxPooling *pool1 = new CUDANet::Layers::MaxPooling(
+        3, 2
+    )
+
+
     return model;
 }
 
diff --git a/tools/export_model_weights.py b/tools/export_model_weights.py
deleted file mode 100644
index 4eb776a..0000000
--- a/tools/export_model_weights.py
+++ /dev/null
@@ -1,28 +0,0 @@
-import torch
-import struct
-
-def export_model_weights(model: torch.nn.Module, filename):
-    with open(filename, 'wb') as f:
-
-        header = ""
-        offset = 0
-
-
-        for name, param in model.named_parameters():
-            if 'weight' not in name and 'bias' not in name:
-                continue
-
-            tensor_values = param.flatten().tolist()
-            tensor_bytes = struct.pack('f' * len(tensor_values), *tensor_values)
-
-            tensor_size = param.numel()
-
-            header += f"{name},{tensor_size},{offset}\n"
-            
-            offset += len(tensor_bytes)
-
-            f.write(tensor_bytes)
-
-        f.seek(0)
-        f.write(struct.pack('q', len(header)))            
-        f.write(header.encode('utf-8'))
diff --git a/tools/utils.py b/tools/utils.py
index 47a53e4..73dbb2c 100644
--- a/tools/utils.py
+++ b/tools/utils.py
@@ -1,7 +1,42 @@
-def print_cpp_vector(vector):
-    print("std::vector<float> expected = {", end="")
+import torch
+import struct
+
+
+def print_cpp_vector(vector, name="expected"):
+    print("std::vector<float> " + name + " = {", end="")
     for i in range(len(vector)):
         if i != 0:
             print(", ", end="")
         print(str(round(vector[i].item(), 5)) + "f", end="")
     print("};")
+
+
+def export_model_weights(model: torch.nn.Module, filename):
+    with open(filename, 'wb') as f:
+
+        header = ""
+        offset = 0
+
+
+        for name, param in model.named_parameters():
+            if 'weight' not in name and 'bias' not in name:
+                continue
+
+            tensor_values = param.flatten().tolist()
+            tensor_bytes = struct.pack('f' * len(tensor_values), *tensor_values)
+
+            tensor_size = param.numel()
+
+            header += f"{name},{tensor_size},{offset}\n"
+            
+            offset += len(tensor_bytes)
+
+            f.write(tensor_bytes)
+
+        f.seek(0)
+        f.write(struct.pack('q', len(header)))            
+        f.write(header.encode('utf-8'))
+
+def print_model_parameters(model: torch.nn.Module):
+    for name, param in model.named_parameters():
+        print(name, param.numel())
\ No newline at end of file