enabled reloading of trained weights + using best weights

app.py

@@ -24,6 +24,8 @@ momentum = 0.5
 log_interval = 10
 random_seed = 1
 TRAIN_CUTOFF = 10
+TEST_PER_SAMPLE = 1500
+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' 
 METRIC_PATH = os.path.join(MODEL_PATH,'metrics.json')
@@ -86,7 +88,7 @@ class MNISTAdversarial_Dataset(Dataset):
         return img, label
 
 class MNISTCorrupted_By_Digit(Dataset):
-    def __init__(self,transform,digit,limit=500):
+    def __init__(self,transform,digit,limit=TEST_PER_SAMPLE):
         self.transform = transform
         self.digit = digit
         corrupted_dir="./mnist_c"
@@ -127,8 +129,8 @@ class MNISTCorrupted(Dataset):
         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'))[:500] for f in files]
-        labels = [np.load(os.path.join(os.path.join(corrupted_dir,f),'test_labels.npy'))[:500] for f in files]
+        images = [np.load(os.path.join(os.path.join(corrupted_dir,f),'test_images.npy'))[:TEST_PER_SAMPLE] for f in files]
+        labels = [np.load(os.path.join(os.path.join(corrupted_dir,f),'test_labels.npy'))[:TEST_PER_SAMPLE] for f in files]
         self.data = np.vstack(images)
         self.labels = np.hstack(labels)
 
@@ -283,24 +285,40 @@ if os.path.exists(model_state_dict) and os.path.exists(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)
+    # 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)
                    
 
-# 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.
+    It loads the latest model weights from the model repository, and then uses those weights to make a
+    prediction on the input image.
     
     :param inp: the image to be classified
-    :return: A dictionary of the class index and the confidence value.
+    :return: A dictionary of the form {class_number: confidence}
     """
+
+    # Get latest model weights ----------------
+    model_repo.git_pull()
+    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)
+    else:
+        # Use 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)
     with torch.no_grad():
 
@@ -314,6 +332,19 @@ def image_classifier(inp):
 
 
 def flag(input_image,correct_result,adversarial_number):
+    """
+    It takes in an image, the correct result, and the number of adversarial images that have been
+    uploaded so far. It saves the image and metadata to a local directory, uploads the image and
+    metadata to the hub, and then pulls the data from the hub to the local directory. If the number of
+    images in the local directory is divisible by the TRAIN_CUTOFF, then it trains the model on the
+    adversarial data
+    
+    :param input_image: The adversarial image that you want to save
+    :param correct_result: The correct number that the image represents
+    :param adversarial_number: This is the number of adversarial examples that have been uploaded to the
+    dataset
+    :return: The output is the output of the flag function.
+    """
 
     adversarial_number = 0 if None else adversarial_number
 
@@ -375,6 +406,12 @@ def get_number_dict(DATA_DIR):
 
 
 def get_statistics():
+    """
+    It loads the model and optimizer state dicts, pulls the latest data from the repo, gets the number
+    of adversarial samples per digit, plots the distribution of adversarial samples per digit, plots the
+    test accuracy per digit per train step, and plots the test accuracy for all digits per train step
+    :return: the following:
+    """
     model_repo.git_pull()
     model_state_dict = MODEL_WEIGHTS_PATH
     optimizer_state_dict = OPTIMIZER_PATH

best_weights/mnist_model.pth

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

best_weights/optimizer.pth

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

data_mnist

@@ -1 +1 @@
-Subproject commit b85a2ad15f628eb33a6595afbaba38cfb6a98ece
+Subproject commit ed62a26e764902f519ff43df850842e07dfe2cc0

utils.py

@@ -24,7 +24,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="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). We are using {TEST_PER_SAMPLE} samples per digit."
 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)."