Project refactor

Utils/models.py

@@ -0,0 +1,51 @@
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.autograd import Variable
+
+
+class MyNetwork(nn.Module): 
+
+	def __init__(self):
+		super(MyNetwork, self).__init__()
+
+		self.conv1 = nn.Conv2d(1, 10, 3)
+		self.pool1 = nn.MaxPool2d(2)
+
+		self.conv2 = nn.Conv2d(10, 20, 3)
+		self.pool2 = nn.MaxPool2d(2)
+
+		self.conv3 = nn.Conv2d(20, 30, 3) 
+		self.dropout = nn.Dropout2d()
+
+		self.fc1 = nn.Linear(30 * 3 * 3, 270) 
+		self.fc2 = nn.Linear(270, 26)
+
+		self.softmax = nn.LogSoftmax(dim=1)
+
+
+	def forward(self, x):
+		x = self.conv1(x)
+		x = F.relu(x)
+		x = self.pool1(x)
+
+		x = self.conv2(x)
+		x = F.relu(x)
+		x = self.pool2(x)
+
+		x = self.conv3(x)
+		x = F.relu(x)
+		x = self.dropout(x)
+
+		x = x.view(-1, 30 * 3 * 3)
+
+		x = self.fc1(x)
+		x = F.relu(x)
+		x = self.fc2(x)
+		x = F.relu(x)
+		x = self.softmax(x)
+
+		return x
+

Utils/utils.py

@@ -0,0 +1,82 @@
+import io
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.autograd import Variable
+from torchvision import transforms
+
+import seaborn as sns
+import torchvision
+from torchsummary import summary
+from torch.utils.tensorboard import SummaryWriter
+
+
+
+def evaluate(predictions, labels):              
+	correct = 0
+	for p, l in zip(predictions, labels):
+		if p == l:
+			correct += 1
+
+	accuracy = correct / len(predictions)
+	return accuracy
+
+
+
+def plot_to_image(figure):
+  """Converts the matplotlib plot specified by 'figure' to a tensor"""
+
+  data = np.frombuffer(figure.canvas.tostring_rgb(), dtype=np.uint8)
+  w, h = figure.canvas.get_width_height()
+  data = data.reshape(h, w, 3)
+
+  trans = transforms.ToPILImage()
+  trans_tensor = transforms.ToTensor()
+  return trans_tensor(trans(data))
+
+def plot_grid(images, labels, predictions, M, N):
+  # Create a figure to contain the plot.
+  figure = plt.figure(figsize=(M,N))
+  for i in range(M*N):
+    # Start next subplot.
+    plt.subplot(M, N, i + 1, title=predictions[i] + ' (' + labels[i] + ')')
+    #plt.setp(title, color=('g' if yp[i*N+j].max(dim=0)[1] == y[i*N+j] else 'r'))
+    plt.xticks([])
+    plt.yticks([])
+    plt.grid(False)
+    plt.imshow(images[i].reshape(28, 28), cmap=plt.cm.binary)
+  plt.show()
+  return figure
+
+def display_pic(pic, position):
+	pixels = pic.reshape(28, 28)
+	plt.subplot(position)
+	sns.heatmap(data=pixels)
+
+def log_image_grid(images, labels, predictions, M, N, writer):
+  plot = plot_grid(images, labels, predictions, M, N)
+  to_show = plot_to_image(plot)
+  writer.add_image('images', to_show, 0)
+
+
+
+def parse_data(path, device):
+	raw_data = pd.read_csv(path, sep=",")
+
+	labels = raw_data['label']
+	raw_data.drop('label', axis=1, inplace=True)
+
+	data = raw_data.values
+	labels = labels.values
+
+	data = data.reshape(data.shape[0], 1, 28, 28)
+
+	x = torch.cuda.FloatTensor(data).to(device)
+	y = torch.cuda.LongTensor(labels).to(device)
+
+	return x, y

attack.py

@@ -0,0 +1,33 @@
+import torch
+
+def fgsm(model, X, y, epsilon):
+    """ Construct FGSM adversarial examples on the examples X"""
+    delta = torch.zeros_like(X, requires_grad=True)
+    loss = nn.CrossEntropyLoss()(model(X + delta), y)
+    loss.backward()
+    return epsilon * delta.grad.detach().sign()
+
+
+def pgd_linf(model, X, y, epsilon, alpha, num_iter):
+    """ Construct PGD adversarial examples on the examples X"""
+    delta = torch.zeros_like(X, requires_grad=True)
+    for t in range(num_iter):
+        loss = nn.CrossEntropyLoss()(model(X + delta), y)
+        loss.backward()
+        delta.data = (delta + alpha*delta.grad.detach().sign()).clamp(-epsilon,epsilon)
+        delta.grad.zero_()
+    return delta.detach()
+
+
+
+def test_attack(model, X, y, attack):
+    delta = attack(model, X, y, 0.1)
+    predictions = model(X + delta)
+
+	#test visualisation
+	M, N = 2, 6
+	images = test_x[0:M*N].detach().cpu()
+	labels = [letters[test_y[i].detach().cpu()] for i in range(M*N)]
+	predictions = [letters[torch.max(predictions.data, 1)[1][i].item()] for i in range(M*N)]
+
+	visualisation.log_image_grid(images, labels, predictions, M, N, writer)

test.py

@@ -0,0 +1,61 @@
+import sys
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+import string
+import random
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.autograd import Variable
+
+import seaborn as sns
+import torchvision
+from torchsummary import summary
+from torch.utils.tensorboard import SummaryWriter
+
+import Utils.utils as utils
+import Utils.models as models
+
+def test(model_path, test_path):
+
+	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+	writer = SummaryWriter()
+
+	model = models.MyNetwork()
+	model.load_state_dict(torch.load(model_path))
+	model = model.to(device)
+
+	test_x, test_y = utils.parse_data(test_path, device)
+
+	letters = string.ascii_uppercase
+	predictions = model(test_x)
+	model.eval()
+	accuracy = utils.evaluate(torch.max(predictions.data, 1)[1], test_y)
+	print("Accuracy: {}".format(accuracy))
+
+	#test visualisation
+	M, N = 2, 6
+	#images = test_x[0:M*N].detach().cpu()
+
+	perm = torch.randperm(test_x.size(0))
+	ids = perm[:M*N]
+
+	images = test_x[ids].detach().cpu()
+
+	labels = [letters[test_y[i].detach().cpu()] for i in ids]
+	predictions = [letters[torch.max(predictions.data, 1)[1][i].item()] for i in ids]
+
+	utils.log_image_grid(images, labels, predictions, M, N, writer)
+
+	writer.close()
+
+if __name__ == "__main__":
+	if len(sys.argv) != 3:
+		print("\nWrong command syntax.\n")
+	else:
+		model_path = sys.argv[1]
+		test_path = sys.argv[2]
+		test(model_path, test_path)

train_regular.py

@@ -0,0 +1,71 @@
+import sys
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.autograd import Variable
+
+import seaborn as sns
+import torchvision
+from torchsummary import summary
+from torch.utils.tensorboard import SummaryWriter
+
+import Utils.utils as utils
+import Utils.models as models
+
+def train_regular(model_path, train_path, num_epochs, batch_size):
+
+	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+	writer = SummaryWriter()
+
+	model = models.MyNetwork().to(device)
+	train_x, train_y = utils.parse_data(train_path, device)
+
+	optimizer = optim.SGD(model.parameters(), 0.001, momentum=0.7)
+	loss_function = nn.CrossEntropyLoss()
+
+	iteration = 0
+	for e in range(num_epochs):
+		for i in range(0, train_x.shape[0], batch_size):
+			train_x_mini = train_x[i:i + batch_size] 
+			train_y_mini = train_y[i:i + batch_size] 
+
+			optimizer.zero_grad()
+			output = model(train_x_mini)
+			loss = loss_function(output, train_y_mini)
+			loss.backward()
+			optimizer.step()
+			predictions = torch.max(output.data, 1)[1]
+
+			if i % 1000 == 0:
+				writer.add_scalar('Loss/Train', loss.item(), iteration)
+				writer.add_scalar('Accuracy/Train', utils.evaluate(predictions, train_y_mini))
+			iteration += 1
+
+		print('Epoch: {} - Loss: {:.6f}'.format(e + 1, loss.item()))
+
+	writer.close()
+	torch.save(model.state_dict(), model_path)
+	summary(model, (1, 28, 28))
+
+
+# python3 train_regular.py <model_path> <train_path> <num_epochs> <batch_size>
+if __name__ == "__main__":
+	if len(sys.argv) != 5:
+		print("\nWrong command syntax.\n")
+	else:
+		good = True
+		try:
+			model_path = sys.argv[1]
+			train_path = sys.argv[2]
+			num_epochs = int(sys.argv[3])
+			batch_size = int(sys.argv[4])
+		except ValueError:
+			print("\nInvalid parameters.\n")
+			good = False
+		if good:
+			train_regular(model_path, train_path, num_epochs, batch_size)