fixing upload of models and metrics

app.py

@@ -231,46 +231,16 @@ optimizer = optim.SGD(network.parameters(), lr=learning_rate,
                       momentum=momentum)        
 
 
-model_state_dict = MODEL_WEIGHTS_PATH
-optimizer_state_dict = OPTIMIZER_PATH
-model_repo.git_pull()
-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)
+def train_and_test(train=True):
 
-    optimizer_state_dict = torch.load(optimizer_state_dict)
-    optimizer.load_state_dict(optimizer_state_dict)                      
+    if train:
+        # Train for one epoch and test
+        train_dataset = MNISTAdversarial_Dataset('./data_mnist',TRAIN_TRANSFORM)
 
-# Train
-#train(n_epochs,network,optimizer)
-
-
-def image_classifier(inp):
-    """
-    It takes an image as input and returns a dictionary of class labels and their corresponding
-    confidence scores.
+        train_loader = torch.utils.data.DataLoader(train_dataset,batch_size=batch_size_test, shuffle=True
+                                                )
+        train(n_epochs,network,optimizer,train_loader)
     
-    :param inp: the image to be classified
-    :return: A dictionary of the class index and the confidence value.
-    """
-    input_image = torchvision.transforms.ToTensor()(inp).unsqueeze(0)
-    with torch.no_grad():
-
-        prediction = torch.nn.functional.softmax(network(input_image)[0], dim=0)
-        #pred_number = prediction.data.max(1, keepdim=True)[1]
-        sorted_prediction = torch.sort(prediction,descending=True)
-        confidences={}
-        for s,v in zip(sorted_prediction.indices.numpy().tolist(),sorted_prediction.values.numpy().tolist()):
-            confidences.update({s:v})
-        return confidences
-
-def train_and_test():
-    # Train for one epoch and test
-    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):
@@ -301,6 +271,48 @@ def train_and_test():
    
     return test_metric
 
+
+model_state_dict = MODEL_WEIGHTS_PATH
+optimizer_state_dict = OPTIMIZER_PATH
+model_repo.git_pull()
+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)        
+                
+else: 
+    # Evaluate model to get initial evaluation with no adversarial training  
+    torch.save(network.state_dict(), MODEL_WEIGHTS_PATH)
+    torch.save(optimizer.state_dict(), OPTIMIZER_PATH)
+    _ = train_and_test(False)
+                   
+
+# Train
+#train(n_epochs,network,optimizer)
+
+
+def image_classifier(inp):
+    """
+    It takes an image as input and returns a dictionary of class labels and their corresponding
+    confidence scores.
+    
+    :param inp: the image to be classified
+    :return: A dictionary of the class index and the confidence value.
+    """
+    input_image = torchvision.transforms.ToTensor()(inp).unsqueeze(0)
+    with torch.no_grad():
+
+        prediction = torch.nn.functional.softmax(network(input_image)[0], dim=0)
+        #pred_number = prediction.data.max(1, keepdim=True)[1]
+        sorted_prediction = torch.sort(prediction,descending=True)
+        confidences={}
+        for s,v in zip(sorted_prediction.indices.numpy().tolist(),sorted_prediction.values.numpy().tolist()):
+            confidences.update({s:v})
+        return confidences
+
+
 def flag(input_image,correct_result,adversarial_number):
 
     adversarial_number = 0 if None else adversarial_number
@@ -380,7 +392,7 @@ def get_statistics():
 
     STATS_EXPLANATION_ = STATS_EXPLANATION.format(num_adv_samples = sum(numbers_count_values))
     
-    plt_digits = plot_bar(numbers_count_values,numbers_count_keys,'Number of adversarial samples',"Digit",f"Distribution of adversarial samples over digits")
+    plt_digits = plot_bar(numbers_count_values,numbers_count_keys,'Number of adversarial samples',"Digit",f"Distribution of adversarial samples per digit")
     
     fig_d, ax_d = plt.subplots(tight_layout=True)
 
@@ -392,7 +404,7 @@ def get_statistics():
                 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={}
 
@@ -404,8 +416,15 @@ def get_statistics():
                 </div>
                """ 
 
-    return plt_digits,fig_d,done_html,STATS_EXPLANATION_    
+    # Plot for total test accuracy for all digits
+    fig_all, ax_all = plt.subplots(tight_layout=True)
+    x_i = [i+1 for i in range(len(metric_dict['all']))]
+
+    ax_all.plot(x_i, metric_dict['all'])
+    fig_all.legend()
+    ax_all.set(xlabel='Adversarial train steps', ylabel='MNIST_C Test Accuracy',title="Test Accuracy for all digits")
 
+    return plt_digits,ax_d.figure,ax_all.figure,done_html,STATS_EXPLANATION_    
 
 
 
@@ -453,8 +472,10 @@ def main():
                 stat_adv_image =gr.Plot(type="matplotlib")
                 gr.Markdown(DASHBOARD_EXPLANATION)
                 test_results=gr.Plot(type="matplotlib")
+                gr.Markdown(DASHBOARD_EXPLANATION_TEST)
+                test_results_all=gr.Plot(type="matplotlib")
 
-            dashboard.select(get_statistics,inputs=[],outputs=[stat_adv_image,test_results,notification,stats])
+            dashboard.select(get_statistics,inputs=[],outputs=[stat_adv_image,test_results,test_results_all,notification,stats])
 
 
 

data_mnist

@@ -1 +1 @@
-Subproject commit fc4a0229b36d92306494d4d1cafb256d1f480046
+Subproject commit b85a2ad15f628eb33a6595afbaba38cfb6a98ece

utils.py

@@ -26,6 +26,7 @@ MODEL_IS_WRONG = """
 DEFAULT_TEST_METRIC = "<html> Current test metric - Avg. loss: 1000, Accuracy: 30/1000 (30%) </html>"
 
 DASHBOARD_EXPLANATION="To test the effect of adversarial training on out-of-distribution data, we track the performance progress of the model on the [MNIST Corrupted test dataset](https://zenodo.org/record/3239543)."
+DASHBOARD_EXPLANATION_TEST="Test accuracy on out-of-distribution data for all numbers."
 
 STATS_EXPLANATION = "Here is the distribution of the __{num_adv_samples}__ adversarial samples we've got. The dataset can be found [here](https://huggingface.co/datasets/chrisjay/mnist-adversarial-dataset)."