Adding new features

mpq/cifar10_dws_cnn/cifar10_dws_cnn.py

@@ -0,0 +1,89 @@
+import init_utils
+import common
+
+# Initialize the environment and get the name
+name = init_utils.initialize_environment(__file__)
+args = init_utils.get_args()
+
+# Set arguments from command line
+max_acc_drop = args.max_acc_drop
+device = args.device
+
+from sklearn.model_selection import train_test_split
+import torch.nn as nn
+import torch.nn.functional as F
+import tensorflow as tf
+import numpy as np
+
+# Load our Dataset
+
+(X_train, y_train), (X_test, y_test) = tf.keras.datasets.cifar10.load_data()
+
+y_train = np.squeeze(y_train, axis = 1)
+y_test = np.squeeze(y_test, axis = 1)
+
+X_train, X_val, y_train, y_val = train_test_split(X_train, y_train, test_size = 0.15)
+
+X_train = (np.transpose(X_train, (0,3,1,2)))
+X_test = (np.transpose(X_test, (0,3,1,2)))
+X_val = (np.transpose(X_val, (0,3,1,2)))
+
+BATCH_SIZE = 128
+epochs = 1
+lr = 0.0001
+
+class DepthwiseBlock(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(DepthwiseBlock, self).__init__()
+        
+        layers = []
+                    
+        layers.append(nn.Conv2d(in_channels = in_channels, out_channels = in_channels, 
+                                    kernel_size = 3, padding = 1, groups = in_channels))  # Depthwise convolution
+        
+        layers.append(nn.ReLU(inplace = True))
+        
+        layers.append(nn.Conv2d(in_channels = in_channels, out_channels = out_channels, 
+                                    kernel_size = 1, padding = 0))  # Pointwise convolution
+            
+        layers.append(nn.ReLU(inplace = True))
+                            
+        self.block = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.block(x)
+    
+class Cifar10_Dws_CNN(nn.Module):
+    def __init__(self):
+        super(Cifar10_Dws_CNN, self).__init__()
+        self.features = nn.Sequential(
+            DepthwiseBlock(in_channels = 3, out_channels = 64),
+            DepthwiseBlock(in_channels = 64, out_channels = 64),
+            nn.MaxPool2d(kernel_size = 2, stride = 2),
+
+            DepthwiseBlock(in_channels = 64, out_channels = 128),
+            DepthwiseBlock(in_channels = 128, out_channels = 128),
+            nn.MaxPool2d(kernel_size = 2, stride = 2),
+            
+            DepthwiseBlock(in_channels = 128, out_channels = 256),
+            DepthwiseBlock(in_channels = 256, out_channels = 256),
+            nn.MaxPool2d(kernel_size = 2, stride = 2)
+        )
+        
+        self.flatten = nn.Flatten()
+        
+        self.classifier = nn.Sequential(
+            nn.Linear(256 * 4 * 4, 10)  # Assuming input size is (32, 32) and after 3 max pooling layers, the size is (4, 4)
+        )
+
+    def forward(self, x):
+        x = self.features(x)
+        x = self.flatten(x)
+        x = self.classifier(x)
+        return F.log_softmax(x, dim = 1)
+
+net = Cifar10_Dws_CNN()
+
+common.create_ibex_qnn(net, name, device, X_train, y_train, X_test, y_test, 
+                X_val = X_val, y_val = y_val, BATCH_SIZE = BATCH_SIZE, 
+                epochs = epochs, lr = lr, max_acc_drop = max_acc_drop)