Export pretrained alexnet

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
+To export weights from pytorch you can use the `export_model_weights` function from `tools/utils.py`  script

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)
+

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;
 }
 

tools/export_model_weights.py

@@ -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'))

tools/utils.py

@@ -1,7 +1,42 @@
-def print_cpp_vector(vector):
+import torch
-    print("std::vector<float> expected = {", end="")
+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())