Add validation

pipelines/train.py

@@ -21,7 +21,7 @@ from cnn import CNNetwork
 # script defaults
 BATCH_SIZE = 128
 EPOCHS = 10
-LEARNING_RATE = 0.001
+LEARNING_RATE = 0.01
 
 TRAIN_FILE="data/train"
 TEST_FILE="data/test"
@@ -48,6 +48,7 @@ def train(
         optimizer,
         origin_device="cuda",
         epochs=10,
+        test_dataloader=None,
     ):
     import os
 
@@ -68,6 +69,8 @@ def train(
     # metrics
     training_acc = []
     training_loss = []
+    testing_acc = []
+    testing_loss = []
 
     wandb.init(project="void-training")
 
@@ -86,6 +89,18 @@ def train(
         now = time.time()
         print("Training Loss: {:.2f}, Training Accuracy: {:.4f}, Time: {:.2f}s".format(training_loss[i], training_acc[i], now - then))
 
+        if test_dataloader:
+            # test model
+            test_epoch_loss, test_epoch_acc = validate_epoch.call(model, test_dataloader, loss_fn, modal_device)
+            
+            # testing metrics
+            testing_loss.append(test_epoch_loss/len(test_dataloader))
+            testing_acc.append(test_epoch_acc/len(test_dataloader))
+
+            print("Testing Loss: {:.2f}, Testing Accuracy  {:.2f}".format(testing_loss[i], testing_acc[i]))
+
+            wandb.log({'testing_loss': testing_loss[i], 'testing_acc': testing_acc[i]})
+
         print ("-------------------------------------------------------- \n")
     
     end = time.time()
@@ -96,7 +111,7 @@ def train(
     return model.to(origin_device)
 
 @stub.function(
-    gpu="any",
+    gpu=gpu.A100(memory=20),
     mounts=[
         Mount.from_local_file(local_path='dataset.py'),
         Mount.from_local_file(local_path='cnn.py'),
@@ -132,6 +147,36 @@ def train_epoch(model, train_dataloader, loss_fn, optimizer, device):
        
     return model, train_loss, train_acc
 
+@stub.function(
+    gpu="any",
+    mounts=[
+        Mount.from_local_file(local_path='dataset.py'),
+        Mount.from_local_file(local_path='cnn.py'),
+    ],
+)
+def validate_epoch(model, test_dataloader, loss_fn, device):
+    from tqdm import tqdm
+
+    test_loss = 0.0
+    test_acc = 0.0
+    total = 0.0
+
+    model.eval()
+
+    with torch.no_grad():
+        for wav, target in tqdm(test_dataloader, "Testing batch..."):
+            wav, target = wav.to(device), target.to(device)
+
+            output = model(wav)
+            loss = loss_fn(output, target)
+
+            test_loss += loss.item()
+            prediciton = torch.argmax(output, 1)
+            test_acc += (prediciton == target).sum().item()/len(prediciton)
+            total += 1
+    
+    return test_loss, test_acc
+
 def save_model(model):
     now = time.strftime("%Y%m%d_%H%M%S")
     model_filename = f"models/void_{now}.pth"
@@ -163,6 +208,9 @@ def main():
     train_dataset = VoiceDataset(TRAIN_FILE, mel_spectrogram, device, time_limit_in_secs=3)
     train_dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
 
+    test_dataset = VoiceDataset(TEST_FILE, mel_spectrogram, device, time_limit_in_secs=3)
+    test_dataloader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=True)
+
     # construct model
     model = CNNetwork()
 
@@ -171,7 +219,7 @@ def main():
     optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
 
     # train model
-    model = train.call(model, train_dataloader, loss_fn, optimizer, device, 3)
+    model = train.call(model, train_dataloader, loss_fn, optimizer, device, EPOCHS, test_dataloader)
 
     # save model
     save_model(model)