diff --git a/.gitattributes b/.gitattributes
index 2c91e8510f64988b36eacc6b4b64244ce27799ea..80afa15e7b516c61723af0bbdc66333499256bb6 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -26,3 +26,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
*.zstandard filter=lfs diff=lfs merge=lfs -text
*tfevents* filter=lfs diff=lfs merge=lfs -text
*ubyte* filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
diff --git a/app.py b/app.py
index 51ea7d02a7b7b01a674d2eaae5f82ef8d06c3982..8230c595559834407fe3688b5d07c9320cc6d765 100644
--- a/app.py
+++ b/app.py
@@ -2,51 +2,163 @@ import os
import torch
import gradio as gr
import torchvision
+from PIL import Image
from utils import *
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from huggingface_hub import Repository, upload_file
+from torch.utils.data import Dataset
+import numpy as np
+from collections import Counter
-n_epochs = 3
-batch_size_train = 64
+
+n_epochs = 10
+batch_size_train = 128
batch_size_test = 1000
learning_rate = 0.01
momentum = 0.5
log_interval = 10
random_seed = 1
-
+TRAIN_CUTOFF = 5
+WHAT_TO_DO=WHAT_TO_DO.format(num_samples=TRAIN_CUTOFF)
+METRIC_PATH = './metrics.json'
REPOSITORY_DIR = "data"
LOCAL_DIR = 'data_local'
os.makedirs(LOCAL_DIR,exist_ok=True)
+
+
HF_TOKEN = os.getenv("HF_TOKEN")
HF_DATASET ="mnist-adversarial-dataset"
+DATASET_REPO_URL = f"https://huggingface.co/datasets/chrisjay/{HF_DATASET}"
+repo = Repository(
+ local_dir="data_mnist", clone_from=DATASET_REPO_URL, use_auth_token=HF_TOKEN
+)
+repo.git_pull()
torch.backends.cudnn.enabled = False
torch.manual_seed(random_seed)
-train_loader = torch.utils.data.DataLoader(
- torchvision.datasets.MNIST('files/', train=True, download=True,
- transform=torchvision.transforms.Compose([
+
+class MNISTAdversarial_Dataset(Dataset):
+
+ def __init__(self,data_dir,transform):
+ repo.git_pull()
+ self.data_dir = os.path.join(data_dir,'data')
+ self.transform = transform
+ files = [f.name for f in os.scandir(self.data_dir)]
+ self.images = []
+ self.numbers = []
+ for f in files:
+ self.FOLDER = os.path.join(os.path.join(self.data_dir,f))
+
+ metadata_path = os.path.join(self.FOLDER,'metadata.jsonl')
+
+ image_path =os.path.join(self.FOLDER,'image.png')
+ if os.path.exists(image_path) and os.path.exists(metadata_path):
+ img = Image.open(image_path)
+ self.images.append(img)
+ metadata = read_json_lines(metadata_path)
+ self.numbers.append(metadata[0]['correct_number'])
+ assert len(self.images)==len(self.numbers), f"Length of images and numbers must be the same. Got {len(self.images)} for images and {len(self.numbers)} for numbers."
+ def __len__(self):
+ return len(self.images)
+
+ def __getitem__(self,idx):
+ img, label = self.images[idx], self.numbers[idx]
+ img = self.transform(img)
+ return img, label
+
+class MNISTCorrupted_By_Digit(Dataset):
+ def __init__(self,transform,digit,limit=30):
+ self.transform = transform
+ self.digit = digit
+ corrupted_dir="./mnist_c"
+ files = [f.name for f in os.scandir(corrupted_dir)]
+ images = [np.load(os.path.join(os.path.join(corrupted_dir,f),'test_images.npy')) for f in files]
+ labels = [np.load(os.path.join(os.path.join(corrupted_dir,f),'test_labels.npy')) for f in files]
+ self.data = np.vstack(images)
+ self.labels = np.hstack(labels)
+
+ assert (self.data.shape[0] == self.labels.shape[0])
+
+ mask = self.labels == self.digit
+
+ data_masked = self.data[mask]
+ # Just to be on the safe side, ensure limit is more than the minimum
+ limit = min(limit,data_masked.shape[0])
+
+ self.data_for_use = data_masked[:limit]
+ self.labels_for_use = self.labels[mask][:limit]
+ assert (self.data_for_use.shape[0] == self.labels_for_use.shape[0])
+
+ def __len__(self):
+ return len(self.data_for_use)
+ def __getitem__(self,idx):
+ if torch.is_tensor(idx):
+ idx = idx.tolist()
+
+ image = self.data_for_use[idx]
+ label = self.labels_for_use[idx]
+ if self.transform:
+ image_pil = torchvision.transforms.ToPILImage()(image) # Need to transform to PIL before using default transforms
+ image = self.transform(image_pil)
+
+ return image, label
+
+
+
+
+
+class MNISTCorrupted(Dataset):
+ def __init__(self,transform):
+ self.transform = transform
+ corrupted_dir="./mnist_c"
+ files = [f.name for f in os.scandir(corrupted_dir)]
+ images = [np.load(os.path.join(os.path.join(corrupted_dir,f),'test_images.npy')) for f in files]
+ labels = [np.load(os.path.join(os.path.join(corrupted_dir,f),'test_labels.npy')) for f in files]
+ self.data = np.vstack(images)
+ self.labels = np.hstack(labels)
+
+ assert (self.data.shape[0] == self.labels.shape[0])
+
+ def __len__(self):
+ return len(self.data)
+
+ def __getitem__(self, idx):
+ if torch.is_tensor(idx):
+ idx = idx.tolist()
+
+ image = self.data[idx]
+ label = self.labels[idx]
+ if self.transform:
+ image_pil = torchvision.transforms.ToPILImage()(image) # Need to transform to PIL before using default transforms
+ image = self.transform(image_pil)
+
+ return image, label
+
+
+
+TRAIN_TRANSFORM = torchvision.transforms.Compose([
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(
(0.1307,), (0.3081,))
- ])),
+ ])
+
+'''
+train_loader = torch.utils.data.DataLoader(
+ torchvision.datasets.MNIST('files/', train=True, download=True,
+ transform=TRAIN_TRANSFORM),
batch_size=batch_size_train, shuffle=True)
+'''
-test_loader = torch.utils.data.DataLoader(
- torchvision.datasets.MNIST('files/', train=False, download=True,
- transform=torchvision.transforms.Compose([
- torchvision.transforms.ToTensor(),
- torchvision.transforms.Normalize(
- (0.1307,), (0.3081,))
- ])),
- batch_size=batch_size_test, shuffle=True)
+test_loader = torch.utils.data.DataLoader(MNISTCorrupted(TRAIN_TRANSFORM),
+ batch_size=batch_size_test, shuffle=False)
# Source: https://nextjournal.com/gkoehler/pytorch-mnist
@@ -69,7 +181,7 @@ class MNIST_Model(nn.Module):
return F.log_softmax(x)
-def train(epochs,network,optimizer):
+def train(epochs,network,optimizer,train_loader):
train_losses=[]
network.train()
@@ -102,10 +214,11 @@ def test():
correct += pred.eq(target.data.view_as(pred)).sum()
test_loss /= len(test_loader.dataset)
test_losses.append(test_loss)
+ acc = 100. * correct / len(test_loader.dataset)
+ acc = acc.item()
test_metric = '〽Current test metric - Avg. loss: `{:.4f}`, Accuracy: `{}/{}` (`{:.0f}%`)\n'.format(
- test_loss, correct, len(test_loader.dataset),
- 100. * correct / len(test_loader.dataset))
- return test_metric
+ test_loss, correct, len(test_loader.dataset),acc )
+ return test_metric,acc
@@ -156,15 +269,41 @@ def image_classifier(inp):
def train_and_test():
# Train for one epoch and test
- train(1,network,optimizer)
- test_metric = test()
-
-def flag(input_image,correct_result,train):
- # take an image, the wrong result, the correct result.
- # push to dataset.
- # get size of current dataset
-
- # Write audio to file
+ train_dataset = MNISTAdversarial_Dataset('./data_mnist',TRAIN_TRANSFORM)
+
+ train_loader = torch.utils.data.DataLoader(train_dataset,batch_size=batch_size_test, shuffle=True
+ )
+ train(n_epochs,network,optimizer,train_loader)
+ test_metric,test_acc = test()
+
+ if os.path.exists(METRIC_PATH):
+ metric_dict = read_json(METRIC_PATH)
+ metric_dict['all'] = metric_dict['all'] if 'all' in metric_dict else [] + [test_acc]
+ else:
+ metric_dict={}
+ metric_dict['all'] = [test_acc]
+
+ for i in range(10):
+ data_per_digit = MNISTCorrupted_By_Digit(TRAIN_TRANSFORM,i)
+ dataloader_per_digit = torch.utils.data.DataLoader(data_per_digit,batch_size=len(data_per_digit), shuffle=False)
+ data_per_digit, label_per_digit = iter(dataloader_per_digit).next()
+ output = network(data_per_digit)
+ pred = output.data.max(1, keepdim=True)[1]
+ correct = pred.eq(label_per_digit.data.view_as(pred)).sum()
+ acc = 100. * correct / len(data_per_digit)
+ acc=acc.item()
+ if os.path.exists(METRIC_PATH):
+ metric_dict[str(i)].append(acc)
+ else:
+ metric_dict[str(i)] = [acc]
+
+ dump_json(thing=metric_dict,file=METRIC_PATH)
+ return test_metric
+
+def flag(input_image,correct_result,adversarial_number):
+
+ adversarial_number = 0 if None else adversarial_number
+
metadata_name = get_unique_name()
SAVE_FILE_DIR = os.path.join(LOCAL_DIR,metadata_name)
os.makedirs(SAVE_FILE_DIR,exist_ok=True)
@@ -182,7 +321,7 @@ def flag(input_image,correct_result,train):
dump_json(metadata,json_file_path)
- # Simply upload the audio file and metadata using the hub's upload_file
+ # Simply upload the image file and metadata using the hub's upload_file
# Upload the image
repo_image_path = os.path.join(REPOSITORY_DIR,os.path.join(metadata_name,'image.png'))
@@ -201,43 +340,84 @@ def flag(input_image,correct_result,train):
repo_type='dataset',
token=HF_TOKEN
)
+ adversarial_number+=1
+ output = f' ✔ ({adversarial_number}) Successfully saved your adversarial data.
'
+ repo.git_pull()
+ length_of_dataset = len([f for f in os.scandir("./data_mnist/data")])
+ test_metric = f" {DEFAULT_TEST_METRIC} "
+ if length_of_dataset % TRAIN_CUTOFF ==0:
+ test_metric_ = train_and_test()
+ test_metric = f" {test_metric_} "
+ output = f' ✔ ({adversarial_number}) Successfully saved your adversarial data and trained the model on adversarial data!
'
+ return output,test_metric,adversarial_number
+
+def get_number_dict(DATA_DIR):
+ files = [f.name for f in os.scandir(DATA_DIR)]
+ numbers = [read_json_lines(os.path.join(os.path.join(DATA_DIR,f),'metadata.jsonl'))[0]['correct_number'] for f in files]
+ numbers_count = Counter(numbers)
+ numbers_count_keys = list(numbers_count.keys())
+ numbers_count_values = [numbers_count[k] for k in numbers_count_keys]
+ return numbers_count_keys,numbers_count_values
+
+
+
+def get_statistics():
+ model_state_dict = 'model.pth'
+ optimizer_state_dict = 'optmizer.pth'
+
+ if os.path.exists(model_state_dict):
+ network_state_dict = torch.load(model_state_dict)
+ network.load_state_dict(network_state_dict)
+
+ if os.path.exists(optimizer_state_dict):
+ optimizer_state_dict = torch.load(optimizer_state_dict)
+ optimizer.load_state_dict(optimizer_state_dict)
+ repo.git_pull()
+ DATA_DIR = './data_mnist/data'
+ numbers_count_keys,numbers_count_values = get_number_dict(DATA_DIR)
+
+
+ plt_digits = plot_bar(numbers_count_values,numbers_count_keys,'Number of adversarial samples',"Digit",f"Distribution of adversarial samples over digits")
- output = f' ✔ Successfully saved to flagged dataset.
'
- train=True
- if train:
- output = f' ✔ Successfully saved to flagged dataset. Training the model on adversarial data!
'
+ fig_d, ax_d = plt.subplots(figsize=(10,4),tight_layout=True)
- return output,train
+ if os.path.exists(METRIC_PATH):
+ metric_dict = read_json(METRIC_PATH)
+ for i in range(10):
+ try:
+ x_i = [i+1 for i in range(len(metric_dict[str(i)]))]
+ ax_d.plot(x_i, metric_dict[str(i)],label=str(i))
+ except Exception:
+ continue
+ dump_json(thing=metric_dict,file=METRIC_PATH)
+ else:
+ metric_dict={}
+ fig_d.legend()
+ ax_d.set(xlabel='Adversarial train steps', ylabel='MNIST_C Test Accuracy',title="Test Accuracy over digits per train step")
+ done_html = """
+
✅ Statistics loaded successfully!
+
+
⌛ Creating statistics...
+