save models and metrics to hub

.gitignore

@@ -1,4 +1,7 @@
 __pycache__/*
 data_local/*
 flagged/*
-data_mnist/*
+data_mnist/*
+model/*
+model
+data_mnist

app.py

@@ -25,7 +25,10 @@ log_interval = 10
 random_seed = 1
 TRAIN_CUTOFF = 10
 WHAT_TO_DO=WHAT_TO_DO.format(num_samples=TRAIN_CUTOFF)
-METRIC_PATH = './metrics.json'
+MODEL_PATH = 'model' 
+METRIC_PATH = os.path.join(MODEL_PATH,'metrics.json')
+MODEL_WEIGHTS_PATH = os.path.join(MODEL_PATH,'mnist_model.pth')
+OPTIMIZER_PATH = os.path.join(MODEL_PATH,'optimizer.pth')
 REPOSITORY_DIR = "data"
 LOCAL_DIR = 'data_local'
 os.makedirs(LOCAL_DIR,exist_ok=True)
@@ -34,14 +37,21 @@ os.makedirs(LOCAL_DIR,exist_ok=True)
 
 
 HF_TOKEN = os.getenv("HF_TOKEN")
-
+MODEL_REPO = 'mnist-adversarial-model'
 HF_DATASET ="mnist-adversarial-dataset"
 DATASET_REPO_URL = f"https://huggingface.co/datasets/chrisjay/{HF_DATASET}"
+MODEL_REPO_URL = f"https://huggingface.co/datasets/chrisjay/{MODEL_REPO}"
+
 repo = Repository(
     local_dir="data_mnist", clone_from=DATASET_REPO_URL, use_auth_token=HF_TOKEN
 )
 repo.git_pull()
 
+model_repo = Repository(
+    local_dir=MODEL_PATH, clone_from=MODEL_REPO_URL, use_auth_token=HF_TOKEN
+)
+model_repo.git_pull()
+
 torch.backends.cudnn.enabled = False
 torch.manual_seed(random_seed)
 
@@ -76,7 +86,7 @@ class MNISTAdversarial_Dataset(Dataset):
         return img, label
 
 class MNISTCorrupted_By_Digit(Dataset):
-    def __init__(self,transform,digit,limit=300):
+    def __init__(self,transform,digit,limit=500):
         self.transform = transform
         self.digit = digit
         corrupted_dir="./mnist_c"
@@ -112,15 +122,13 @@ class MNISTCorrupted_By_Digit(Dataset):
 
         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'))[:300] for f in files]
-        labels = [np.load(os.path.join(os.path.join(corrupted_dir,f),'test_labels.npy'))[:300] for f in files]
+        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]
         self.data = np.vstack(images)
         self.labels = np.hstack(labels)
 
@@ -142,7 +150,6 @@ class MNISTCorrupted(Dataset):
         return image, label
 
 
-
 TRAIN_TRANSFORM = torchvision.transforms.Compose([
                                torchvision.transforms.ToTensor(),
                                torchvision.transforms.Normalize(
@@ -191,8 +198,8 @@ def train(epochs,network,optimizer,train_loader):
                 100. * batch_idx / len(train_loader), loss.item()))
                 train_losses.append(loss.item())
         
-                torch.save(network.state_dict(), 'model.pth')
-                torch.save(optimizer.state_dict(), 'optimizer.pth')
+                torch.save(network.state_dict(), MODEL_WEIGHTS_PATH)
+                torch.save(optimizer.state_dict(), OPTIMIZER_PATH)
 
 def test():
     test_losses=[]
@@ -224,19 +231,16 @@ optimizer = optim.SGD(network.parameters(), lr=learning_rate,
                       momentum=momentum)        
 
 
-model_state_dict = 'model.pth'
-optimizer_state_dict = 'optmizer.pth'
-
-if os.path.exists(model_state_dict):
+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)
 
-if os.path.exists(optimizer_state_dict):
     optimizer_state_dict = torch.load(optimizer_state_dict)
     optimizer.load_state_dict(optimizer_state_dict)                      
 
-
-
 # Train
 #train(n_epochs,network,optimizer)
 
@@ -291,6 +295,10 @@ def train_and_test():
             metric_dict[str(i)] = [acc] 
 
     dump_json(thing=metric_dict,file=METRIC_PATH)    
+
+    # Push models and metrics to hub
+    model_repo.push_to_hub()
+   
     return test_metric
 
 def flag(input_image,correct_result,adversarial_number):
@@ -355,8 +363,9 @@ def get_number_dict(DATA_DIR):
 
 
 def get_statistics():
-    model_state_dict = 'model.pth'
-    optimizer_state_dict = 'optmizer.pth'
+    model_repo.git_pull()
+    model_state_dict = MODEL_WEIGHTS_PATH
+    optimizer_state_dict = OPTIMIZER_PATH
 
     if os.path.exists(model_state_dict):
         network_state_dict = torch.load(model_state_dict)

data_mnist

@@ -1 +1 @@
-Subproject commit 650e2ac4a86b5e109a12b5adc7bf6436bbe578de
+Subproject commit 5915a9276e314d92a5b533b5312616b28b9bcee5