alexmr09
78 lines
2.3 KiB
Python
78 lines
2.3 KiB
Python
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)) - 128.0)/255.0
|
|
X_test = (np.transpose(X_test, (0,3,1,2)) - 128.0)/255.0
|
|
X_val = (np.transpose(X_val, (0,3,1,2)) - 128.0)/255.0
|
|
|
|
BATCH_SIZE = 32
|
|
epochs = 1
|
|
lr = 0.0001
|
|
|
|
class CMSIS_CNN(nn.Module):
|
|
def __init__(self):
|
|
super(CMSIS_CNN, self).__init__()
|
|
self.conv1 = nn.Conv2d(in_channels = 3, out_channels = 32, kernel_size = 5, padding = 2)
|
|
self.relu1 = nn.ReLU()
|
|
self.max1 = nn.MaxPool2d(2,2)
|
|
self.d1 = nn.Dropout(p = 0.25)
|
|
|
|
self.conv2 = nn.Conv2d(in_channels = 32, out_channels = 32, kernel_size = 5, padding = 2)
|
|
self.relu2 = nn.ReLU()
|
|
self.max2 = nn.MaxPool2d(2,2)
|
|
self.d2 = nn.Dropout(p = 0.25)
|
|
|
|
self.conv3 = nn.Conv2d(in_channels = 32, out_channels = 64, kernel_size = 5, padding = 2)
|
|
self.relu3 = nn.ReLU()
|
|
self.max3 = nn.MaxPool2d(2,2)
|
|
self.d3 = nn.Dropout(p = 0.4)
|
|
|
|
self.flatten = nn.Flatten()
|
|
self.linear1 = nn.Linear(1024, 10)
|
|
|
|
def forward(self,X):
|
|
X = self.relu1((self.conv1(X)))
|
|
X = self.max1(X)
|
|
X = self.d1(X)
|
|
|
|
X = self.relu2((self.conv2(X)))
|
|
X = self.max2(X)
|
|
X = self.d2(X)
|
|
|
|
X = self.relu3((self.conv3(X)))
|
|
X = self.max3(X)
|
|
X = self.d3(X)
|
|
|
|
X = self.flatten(X)
|
|
|
|
X = self.linear1(X)
|
|
return F.log_softmax(X, dim = 1)
|
|
|
|
net = CMSIS_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)
|