Merge branch 'main' of https://huggingface.co/spaces/chrisjay/mnist-adversarial

.gitignore

@@ -4,4 +4,5 @@ flagged/*
 data_mnist/*
 model/*
 model
-data_mnist
+data_mnist
+slurm*

README.md

@@ -10,3 +10,4 @@ pinned: false
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+

app.py

@@ -20,11 +20,12 @@ n_epochs = 10
 batch_size_train = 128
 batch_size_test = 1000
 learning_rate = 0.01
+adv_learning_rate= 0.001
 momentum = 0.5
 log_interval = 10
 random_seed = 1
 TRAIN_CUTOFF = 10
-TEST_PER_SAMPLE = 1500
+TEST_PER_SAMPLE = 5000
 DASHBOARD_EXPLANATION = DASHBOARD_EXPLANATION.format(TEST_PER_SAMPLE=TEST_PER_SAMPLE)
 WHAT_TO_DO=WHAT_TO_DO.format(num_samples=TRAIN_CUTOFF)
 MODEL_PATH = 'model' 
@@ -163,7 +164,6 @@ TRAIN_TRANSFORM = torchvision.transforms.Compose([
 test_loader = torch.utils.data.DataLoader(MNISTCorrupted(TRAIN_TRANSFORM),
                                             batch_size=batch_size_test, shuffle=False)
 
-
 # Source: https://nextjournal.com/gkoehler/pytorch-mnist
 class MNIST_Model(nn.Module):
     def __init__(self):
@@ -221,6 +221,7 @@ def test():
         acc = acc.item()
         test_metric = '〽Current test metric -> Avg. loss: `{:.4f}`, Accuracy: `{:.0f}%`\n'.format(
         test_loss,acc)
+        print(test_metric)
         return test_metric,acc
 
 
@@ -234,6 +235,34 @@ optimizer = optim.SGD(network.parameters(), lr=learning_rate,
                       momentum=momentum)        
 
 
+
+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_iid_loader = torch.utils.data.DataLoader(
+  torchvision.datasets.MNIST('./files/', train=False, download=True,
+                             transform=TRAIN_TRANSFORM),
+  batch_size=batch_size_test, shuffle=True)
+
+model_state_dict = MODEL_WEIGHTS_PATH
+optimizer_state_dict = OPTIMIZER_PATH
+if os.path.exists(model_state_dict) and os.path.exists(optimizer_state_dict):
+    network_state_dict = torch.load(model_state_dict)
+    network.load_state_dict(network_state_dict)
+
+    optimizer_state_dict = torch.load(optimizer_state_dict)
+    optimizer.load_state_dict(optimizer_state_dict)     
+
+# Train model
+#n_epochs=20
+#train(n_epochs,network,optimizer,train_loader)
+#test()
+
+
+
+
 def train_and_test(train_model=True):
 
     if train_model:
@@ -245,6 +274,7 @@ def train_and_test(train_model=True):
     
     test_metric,test_acc = test()
 
+    network.eval()
     if os.path.exists(METRIC_PATH):
         metric_dict = read_json(METRIC_PATH)
         metric_dict['all'] = metric_dict['all']+ [test_acc] if 'all' in metric_dict else [] + [test_acc]
@@ -274,6 +304,7 @@ def train_and_test(train_model=True):
     return test_metric
 
 
+# Update model weights again
 model_state_dict = MODEL_WEIGHTS_PATH
 optimizer_state_dict = OPTIMIZER_PATH
 model_repo.git_pull()
@@ -288,9 +319,14 @@ else:
     # Use best weights  
     BEST_WEIGHTS_MODEL = "best_weights/mnist_model.pth"
     BEST_WEIGHTS_OPTIMIZER = "best_weights/optimizer.pth"
-    torch.save(network.state_dict(), BEST_WEIGHTS_MODEL)
-    torch.save(optimizer.state_dict(), BEST_WEIGHTS_OPTIMIZER)
-    _ = train_and_test(False)
+    
+    network_state_dict = torch.load(BEST_WEIGHTS_MODEL)
+    network.load_state_dict(network_state_dict)
+
+    optimizer_state_dict = torch.load(BEST_WEIGHTS_OPTIMIZER)
+    optimizer.load_state_dict(optimizer_state_dict)        
+    if not os.path.exists(METRIC_PATH):
+        _ = train_and_test(False)
                    
 
 def image_classifier(inp):
@@ -306,20 +342,23 @@ def image_classifier(inp):
     model_repo.git_pull()
     model_state_dict = MODEL_WEIGHTS_PATH
     optimizer_state_dict = OPTIMIZER_PATH
+    which_weights=''
 
     if os.path.exists(model_state_dict) and os.path.exists(optimizer_state_dict):
+        which_weights = "Using weights from model repo"
         network_state_dict = torch.load(model_state_dict)
         network.load_state_dict(network_state_dict)
         optimizer_state_dict = torch.load(optimizer_state_dict)
         optimizer.load_state_dict(optimizer_state_dict)
     else:
-        # Use best weights  
+        # Use best weights 
+        which_weights = "Using default best weights"
         BEST_WEIGHTS_MODEL = "best_weights/mnist_model.pth"
         BEST_WEIGHTS_OPTIMIZER = "best_weights/optimizer.pth"   
         network.load_state_dict(torch.load(BEST_WEIGHTS_MODEL))
         optimizer.load_state_dict(torch.load(BEST_WEIGHTS_OPTIMIZER)) 
-
-    input_image = torchvision.transforms.ToTensor()(inp).unsqueeze(0)
+    network.eval()
+    input_image = TRAIN_TRANSFORM(inp).unsqueeze(0)
     with torch.no_grad():
 
         prediction = torch.nn.functional.softmax(network(input_image)[0], dim=0)

best_weights/mnist_model.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ba8d282674beb300db53069e4972cfed358f8c7c627cf449215e44b365fcdc54
+oid sha256:148112958ca9545938f0660cec604ac4c7f52dca3523091e1e8e4e6a26e1ebc7
 size 89871

best_weights/optimizer.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:fe255c0ca501d01ae3c2083ea760ea95759fcfe9075e39fba299c57a9907bf1b
-size 623
+oid sha256:aac1c136737d50c2665563392a5b220396398cf1e2a2049dbefd7dc95473f5a5
+size 89807

requirements.txt

@@ -1,3 +1,5 @@
 torch
 torchvision
-matplotlib
+matplotlib
+gradio
+huggingface_hub