added avg validation

Yolo_train.py

@@ -31,44 +31,14 @@ transform = transforms.Compose([
 ])
 
 
-SAVE_EVERY = 2
-batch_size = 32
-
-trainSet = CocoDataSet("./data/train2014", "./data/labels/train2014", transform, 1) 
-trainLoader = DataLoader(trainSet, batch_size=batch_size,
-                                          shuffle=True, num_workers=0)
-
-valSet = CocoDataSet("./data/val2014", "./data/labels/val2014", transform, 1) 
-valLoader = DataLoader(valSet, batch_size=batch_size,
-                                          shuffle=True, num_workers=0)
-
-yoloNetwork = Yolo_v1_fcs(1280*7*7)
-
-
-# CUDA setup guide - https://www.youtube.com/watch?v=GMSjDTU8Zlc
-device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
-# convert model to cuda if available
-if torch.cuda.is_available():
-    encoder.to('cuda')
-    yoloNetwork.fcs.to('cuda')
-
-#Loss stuff:
-criterion = YoloLoss()
-optimizer = optim.Adam(yoloNetwork.parameters(), lr=0.0005)#, momentum=0.9)
-
-if (len(sys.argv) > 1):
-    if (sys.argv[1] == "--load-model"):
-        save_file = sys.argv[2]
-        utils.load_YoloModel_and_optimizer(yoloNetwork, optimizer, save_file)
-        # model, optimizer = utils.load_model_and_optimizer(save_file)
-
-for epoch in range(2):
+def train_one_epoch(yolo_network, train_loader):
     running_loss = []
     
-    for i, data in enumerate(trainLoader):
+    for i, data in enumerate(train_loader):
         # get the inputs; data is a list of [inputs, labels]
-
         
+        if (i == 50): break
+
         inputs, labels = data[0].to(device), data[1].to(device)
 
         # zero the parameter gradients
@@ -76,7 +46,7 @@ for epoch in range(2):
 
         # forward + backward + optimize
         inputs = encoder(inputs)
-        outputs = yoloNetwork(inputs)
+        outputs = yolo_network(inputs)
         loss = criterion(outputs, labels)
         loss.backward()
         optimizer.step()
@@ -85,14 +55,79 @@ for epoch in range(2):
         running_loss.append(loss.item())
 
         # if epoch % 3 == 0:
-        if i % 5 == 0 and not i == 0:
+        if i % 25 == 0 and not i == 0:
             print(f'[{epoch + 1}, {i + 1:5d}] loss: {sum(running_loss) / (i+1)}')
-            checkpoint_state = {"state_dict": yoloNetwork.fcs.state_dict(), "optimizer": optimizer.state_dict()}
-            save_name = utils.save_model(checkpoint_state)
-            print(f"Saved model as '{save_name}'", end="\n\n")
+            # utils.save_model(yolo_network, optimizer)
 
         if i % 2000 == 1999: # print every 2000 mini-batches
             print(f'[{epoch + 1}, {i + 1:5d}] loss: {sum(running_loss) / 2000:.3f}')
             running_loss = 0.0
+    
+
+if __name__ == '__main__':
+    
+    EPOCHS = 5
+    BATCH_SIZE = 32
+    best_avg_val_loss = float("inf")
+
+    print("creating train dataset and loader")
+    train_set = CocoDataSet("./data/train2014", "./data/labels/train2014", transform, 1) 
+    train_loader = DataLoader(train_set, batch_size=BATCH_SIZE,
+                                            shuffle=True, num_workers=0)
+
+    print("creating val dataset and loader")
+    val_set = CocoDataSet("./data/val2014", "./data/labels/val2014", transform, 1) 
+    val_loader = DataLoader(val_set, batch_size=BATCH_SIZE,
+                                            shuffle=True, num_workers=0)
+
+    yolo_network = Yolo_v1_fcs(1280*7*7)
+
+    # CUDA setup guide - https://www.youtube.com/watch?v=GMSjDTU8Zlc
+    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+    # convert model to cuda if available
+    if torch.cuda.is_available():
+        encoder.to('cuda')
+        yolo_network.fcs.to('cuda')
+
+    #Loss stuff:
+    criterion = YoloLoss()
+    optimizer = optim.Adam(yolo_network.parameters(), lr=0.001)#, momentum=0.9)
+
+    if (len(sys.argv) > 1):
+        if (sys.argv[1] == "--load-model"):
+            save_file = sys.argv[2]
+            utils.load_YoloModel_and_optimizer(yolo_network, optimizer, save_file)
+            # model, optimizer = utils.load_model_and_optimizer(save_file)
+
+    for epoch in range(EPOCHS):
+        print(f"__Epoch {epoch+1}__")
+        
+        yolo_network.train(True)
+        train_one_epoch(yolo_network, train_loader)
+
+        yolo_network.train(False)
+        
+        running_val_loss = 0
+        for i, val_data in enumerate(val_loader):
+            if (i == 25): break
+            
+            inputs, labels = val_data[0].to(device), val_data[1].to(device)
+
+            inputs = encoder(inputs)
+            outputs = yolo_network(inputs)
+
+            val_loss = criterion(outputs, labels)
+            val_loss.backward()
+
+            running_val_loss += val_loss
+
+        avg_val_loss = running_val_loss/(i+1)
+
+        if (avg_val_loss < best_avg_val_loss):
+            print(f"New best avg. validation loss: {avg_val_loss}")
+            best_avg_val_loss = avg_val_loss
+            utils.save_model(yolo_network, optimizer)
+
+    
 
-print('Finished Training')
\ No newline at end of file
+    print('Finished Training')
\ No newline at end of file

utils.py

@@ -9,7 +9,7 @@ def predicted_bb_to_draw(b):
     """
     return (b[0] - b[2]/2,b[1] - b[3]/2, *b[2:] )
 
-def save_model(state, file_name=None, folder=fr"saved_models\{datetime.strftime(datetime.now(), '%Y_%m_%d')}"):
+def save_model(yolo_network, optimizer, file_name=None, folder=fr"saved_models\{datetime.strftime(datetime.now(), '%Y_%m_%d')}"):
     """
     Save the model
     based on https://www.youtube.com/watch?v=g6kQl_EFn84
@@ -18,11 +18,13 @@ def save_model(state, file_name=None, folder=fr"saved_models\{datetime.strftime(
     file name should be .pth.tar
     if no file name is given it will default to the time of the save in the format "hour_minute_second"
     """
-
+    state = {"state_dict": yolo_network.fcs.state_dict(), "optimizer": optimizer.state_dict()}
+    
     if file_name is None:
         file_name = fr"{folder}\{datetime.strftime(datetime.now(), '%H_%M_%S')}.pth.tar"
     torch.save(state, file_name)
-    return file_name
+
+    print(f"Saved model as '{file_name}'", end="\n\n")
 
 def load_YoloModel_and_optimizer(YoloModel, optimizer, load_file, folder=fr"saved_models\{datetime.strftime(datetime.now(), '%Y_%m_%d')}"):
     print(f"___Loading model and optimizer states from {load_file}___")