Add min configuration for this.

config.yaml

@@ -6,8 +6,8 @@ model:
   epochs: 10
   batch_size: 8
   optimizer: adam
-  image_sample_size: 10
-  trainable_backbone_layers: 3
+  image_sample_size: 10000
+  trainable_backbone_layers: 1
 
 dataset:
   image_dir: datasets/images
@@ -16,10 +16,10 @@ dataset:
   val_split: 0.2
 
 notebooks:
-  visualization: models/training_log.json
+  visualization: logs/training_log.json
 
 logs:
-  log_dir: models/training_log.txt
+  log_dir: logs/training_log.txt
   log_interval: 10
 
 evaluation:

src/train.py

@@ -59,6 +59,8 @@ def train_one_epoch(model, optimizer, data_loader, device):
 
         loss_dict = model(images, targets)
         loss = sum(loss for loss in loss_dict.values())
+        if(batch_idx % 100 == 0):
+            log_message(f"Iteration {batch_idx+1}, Loss: {loss.item()}")
 
         optimizer.zero_grad()
         loss.backward()
@@ -70,7 +72,7 @@ def train_one_epoch(model, optimizer, data_loader, device):
 
 def store_log(loss, mAP):
     # Save log in JSON for visualization
-    log_json = config["logs"]["log_dir"]
+    log_json = config["notebooks"]["visualization"]
     if not os.path.exists(log_json):
         log_data = {"loss": [], "mAP": []}
     else: