Adding new features

2024-07-23 13:00:49 +03:00
parent 9e044fd7fc
commit 745cc4ed6d
28 changed files with 33632 additions and 106 deletions
@@ -151,15 +151,154 @@ class Ibex_Lenet5(nn.Module):

        return X

+class Ibex_CMSIS_CNN(nn.Module):
+    def __init__(self, mul_vals, shift_vals):
+        super(Ibex_CMSIS_CNN, self).__init__()
+        
+        self.m0 = mul_vals[0]
+        self.m1 = mul_vals[1]
+        self.m2 = mul_vals[2]
+        self.m3 = mul_vals[3]
+        
+        self.s0 = shift_vals[0] + 7
+        self.s1 = shift_vals[1] + 7
+        self.s2 = shift_vals[2] + 7
+        self.s3 = shift_vals[3] + 7
+        
+        self.conv1 = nn.Conv2d(in_channels = 3, out_channels = 32, kernel_size = 5, padding = 'same')
+        self.max1 = nn.MaxPool2d(2,2)
+        
+        self.conv2 = nn.Conv2d(in_channels = 32, out_channels = 32, kernel_size = 5, padding = 'same')
+        self.max2 = nn.MaxPool2d(2,2)
+        
+        self.conv3 = nn.Conv2d(in_channels = 32, out_channels = 64, kernel_size = 5, padding = 'same')
+        self.max3 = nn.MaxPool2d(2,2)
+        
+        self.linear1 = nn.Linear(1024, 10)
+        
+    def forward(self, X, print_out = False):
+        
+        X = self.conv1(X)
+        X = torch.mul(X, self.m0)
+        X = torch.add(X, torch.bitwise_left_shift(torch.tensor(1), self.s0 -1)).type(torch.LongTensor)
+        X = torch.bitwise_right_shift(X, self.s0).type(torch.FloatTensor)
+        X = torch.clamp(X, min = 0, max = 255)
+        
+        X = self.max1(X)
+        
+        X = self.conv2(X)
+        X = torch.mul(X, self.m1)
+        X = torch.add(X, torch.bitwise_left_shift(torch.tensor(1), self.s1 -1)).type(torch.LongTensor)
+        X = torch.bitwise_right_shift(X, self.s1).type(torch.FloatTensor)
+        X = torch.clamp(X, min = 0, max = 255)
+        
+        X = self.max2(X)
+        
+        X = self.conv3(X)
+        X = torch.mul(X, self.m2)
+        X = torch.add(X, torch.bitwise_left_shift(torch.tensor(1), self.s2 -1)).type(torch.LongTensor)
+        X = torch.bitwise_right_shift(X, self.s2).type(torch.FloatTensor)
+        X = torch.clamp(X, min = 0, max = 255)
+        
+        X = self.max3(X)
+        
+        X = X.reshape(X.shape[0], -1)
+        X = self.linear1(X)
+        X = torch.mul(X, self.m3)
+        X = torch.add(X, torch.bitwise_left_shift(torch.tensor(1), self.s3 -1)).type(torch.LongTensor)
+        X = torch.bitwise_right_shift(X, self.s3).type(torch.FloatTensor)
+        X = torch.clamp(X, min = 0, max = 255)
+
+        if(print_out):
+            print(X)
+        return X
+
+class Ibex_DepthwiseBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, mul_vals, shift_vals):
+        super(Ibex_DepthwiseBlock, self).__init__()
+                            
+        self.dw = nn.Conv2d(in_channels = in_channels, out_channels = in_channels, 
+                                    kernel_size = 3, padding = 1, groups = in_channels)
+                
+        self.pw = nn.Conv2d(in_channels = in_channels, out_channels = out_channels, 
+                                    kernel_size = 1, padding = 0)
+            
+        self.m0 = mul_vals[0]
+        self.m1 = mul_vals[1]
+        
+        self.s0 = shift_vals[0] + 7
+        self.s1 = shift_vals[1] + 7
+        
+    def forward(self, X):
+        X = self.dw(X)
+        X = torch.mul(X, self.m0)
+        X = torch.add(X, torch.bitwise_left_shift(torch.tensor(1), self.s0 -1)).type(torch.LongTensor)
+        X = torch.bitwise_right_shift(X, self.s0).type(torch.FloatTensor)
+        X = torch.clamp(X, min = 0, max = 255)
+        
+        X = self.pw(X)
+        X = torch.mul(X, self.m1)
+        X = torch.add(X, torch.bitwise_left_shift(torch.tensor(1), self.s1 -1)).type(torch.LongTensor)
+        X = torch.bitwise_right_shift(X, self.s1).type(torch.FloatTensor)
+        X = torch.clamp(X, min = 0, max = 255)
+        
+        return X
+    
+class Ibex_Cifar10_Dws_CNN(nn.Module):
+    def __init__(self, mul_vals, shift_vals):
+        super(Ibex_Cifar10_Dws_CNN, self).__init__()
+        self.features = nn.Sequential(
+            Ibex_DepthwiseBlock(3, 64, mul_vals[0:2], shift_vals[0:2]),
+            Ibex_DepthwiseBlock(64, 64, mul_vals[2:4], shift_vals[2:4]),
+            nn.MaxPool2d(kernel_size = 2, stride = 2),
+
+            Ibex_DepthwiseBlock(64, 128, mul_vals[4:6], shift_vals[4:6]),
+            Ibex_DepthwiseBlock(128, 128, mul_vals[6:8], shift_vals[6:8]),
+            nn.MaxPool2d(kernel_size = 2, stride = 2),
+            
+            Ibex_DepthwiseBlock(128, 256, mul_vals[8:10], shift_vals[8:10]),
+            Ibex_DepthwiseBlock(256, 256, mul_vals[10:12], shift_vals[10:12]),
+            nn.MaxPool2d(kernel_size = 2, stride = 2)
+        )
+        
+        self.flatten = nn.Flatten()
+        
+        self.classifier = nn.Sequential(
+            nn.Linear(256 * 4 * 4, 10)
+        )
+        
+        self.m_cl = mul_vals[12]
+        self.s_cl = shift_vals[12] + 7
+
+    def forward(self, x, print_out = False):
+        x = self.features(x)
+        x = self.flatten(x)
+        x = self.classifier(x)
+        
+        x = torch.mul(x, self.m_cl)
+        x = torch.add(x, torch.bitwise_left_shift(torch.tensor(1), self.s_cl - 1)).type(torch.LongTensor)
+        x = torch.bitwise_right_shift(x, self.s_cl).type(torch.FloatTensor)
+        x = torch.clamp(x, min = 0, max = 255)
+        
+        if(print_out):
+            print(x)
+            
+        return x
+
+def configure_network(ibex_model_dict, int_weights, int_biases):
+    for i, (name, _) in enumerate(ibex_model_dict.items()):
+        if(i%2 == 0):
+            ibex_model_dict[name] =  torch.tensor(int_weights[i//2])
+        else:
+            ibex_model_dict[name] = torch.tensor(int_biases[i//2])
+
+    return ibex_model_dict
+
 def create_fann_model(int_weights, int_biases, mul_vals, shift_vals):
    ibex_model = Ibex_FANN(mul_vals, shift_vals)
    ibex_model_dict = ibex_model.state_dict()

-    ibex_model_dict['linear1.weight'] = torch.tensor(int_weights[0])
-    ibex_model_dict['linear2.weight'] = torch.tensor(int_weights[1])
-
-    ibex_model_dict['linear1.bias'] = torch.tensor(int_biases[0])
-    ibex_model_dict['linear2.bias'] = torch.tensor(int_biases[1])
+    ibex_model_dict = configure_network(ibex_model_dict, int_weights, int_biases)

    ibex_model.load_state_dict(ibex_model_dict)
    return ibex_model
@@ -168,15 +307,7 @@ def create_uci_model(int_weights, int_biases, mul_vals, shift_vals):
    ibex_model = Ibex_UCI_MLP(mul_vals, shift_vals)
    ibex_model_dict = ibex_model.state_dict()

-    ibex_model_dict['fc0.weight'] = torch.tensor(int_weights[0])
-    ibex_model_dict['fc1.weight'] = torch.tensor(int_weights[1])
-    ibex_model_dict['fc2.weight'] = torch.tensor(int_weights[2])
-    ibex_model_dict['fc3.weight'] = torch.tensor(int_weights[3])
-
-    ibex_model_dict['fc0.bias'] = torch.tensor(int_biases[0])
-    ibex_model_dict['fc1.bias'] = torch.tensor(int_biases[1])
-    ibex_model_dict['fc2.bias'] = torch.tensor(int_biases[2])
-    ibex_model_dict['fc3.bias'] = torch.tensor(int_biases[3])
+    ibex_model_dict = configure_network(ibex_model_dict, int_weights, int_biases)

    ibex_model.load_state_dict(ibex_model_dict)

@@ -186,24 +317,33 @@ def create_lenet_model(int_weights, int_biases, mul_vals, shift_vals):
    ibex_model = Ibex_Lenet5(mul_vals, shift_vals)
    ibex_model_dict = ibex_model.state_dict()

-    ibex_model_dict['conv1.weight'] = torch.tensor(int_weights[0])
-    ibex_model_dict['conv2.weight'] = torch.tensor(int_weights[1])
-    ibex_model_dict['fc1.weight'] = torch.tensor(int_weights[2])
-    ibex_model_dict['fc2.weight'] = torch.tensor(int_weights[3])
-    ibex_model_dict['fc3.weight'] = torch.tensor(int_weights[4])
-
-    ibex_model_dict['conv1.bias'] = torch.tensor(int_biases[0])
-    ibex_model_dict['conv2.bias'] = torch.tensor(int_biases[1])
-    ibex_model_dict['fc1.bias'] = torch.tensor(int_biases[2])
-    ibex_model_dict['fc2.bias'] = torch.tensor(int_biases[3])
-    ibex_model_dict['fc3.bias'] = torch.tensor(int_biases[4])
+    ibex_model_dict = configure_network(ibex_model_dict, int_weights, int_biases)

    ibex_model.load_state_dict(ibex_model_dict)

    return ibex_model

+def create_cmsis_cnn_model(int_weights, int_biases, mul_vals, shift_vals):
+    ibex_model = Ibex_CMSIS_CNN(mul_vals, shift_vals)
+    ibex_model_dict = ibex_model.state_dict()
+
+    ibex_model_dict = configure_network(ibex_model_dict, int_weights, int_biases)
+
+    ibex_model.load_state_dict(ibex_model_dict)
+
+    return ibex_model
+
+def create_ibex_dws_model(int_weights, int_biases, mul_vals, shift_vals):
+    ibex_model = Ibex_Cifar10_Dws_CNN(mul_vals, shift_vals)
+    ibex_model_dict = ibex_model.state_dict()
+
+    ibex_model_dict = configure_network(ibex_model_dict, int_weights, int_biases)
+    
+    ibex_model.load_state_dict(ibex_model_dict)
+    
+    return ibex_model
+
 def eval_sim_model(quant_model, ibex_model, test_loader):
-    # Turn off gradients for validation
    with torch.no_grad():
        ibex_model.eval()
        correct = 0