update

main.py

@@ -6,17 +6,15 @@ import torch.nn.functional as F
 from src.models.model import ShapeClassifier
 
 from src.configs.model_config import ModelConfig
-from src.data.data_loader import train_loader, num_classes
-from src.utils.tensorboard import writer
+from src.data.data_loader import train_loader, num_classes, val_loader
 from src.utils.train import train
 from src.utils.test import test
 from src.utils.wandb import wandb
 from src.utils.logs import logging
-import json
 from src.utils.model import save_model
 
 
-def train_runner():
+def main():
 
     config = ModelConfig().get_config()
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -30,7 +28,7 @@ def train_runner():
 
         loss = train(train_loader, model=model, loss_fn=F.cross_entropy,
                      optimizer=optimizer)
-        test(train_loader, model=model, loss_fn=F.cross_entropy)
+        test(val_loader, model=model, loss_fn=F.cross_entropy)
         # 3. Log metrics over time to visualize performance
         wandb.log({"loss": loss})
 
@@ -38,10 +36,10 @@ def train_runner():
         save_model(model, "results/models/last.pth")
 
         # 4. Log an artifact to W&B
-        wandb.log_artifact("model.pth")
+        # wandb.log_artifact("model.pth")
         # model.train()
 
 
 if __name__ == "__main__":
     logging.info("Training model")
-    train_runner()
+    main()

src/data/data_loader.py

@@ -3,26 +3,35 @@ from torch.utils.data import DataLoader
 from src.configs.model_config import ModelConfig
 from .transform import data_transform
 import os
+from torch.utils.data import random_split
+
 
 num_classes = 3
 config = ModelConfig().get_config()
 
-train_dataset = CustomDataset(data_folder=os.path.join(
+all_dataset = CustomDataset(data_folder=os.path.join(
     "data", 'raw'), transform=data_transform)
 
-# # Calculate the split point
-# split_index = int(0.8 * len(dataset))
+# split to train, val, test
+total_size = len(all_dataset)
+train_size = int(0.8 * total_size)
+val_size = int(0.1 * total_size)
+test_size = total_size - train_size - val_size
 
-# # Split the dataset into training and testing
-# train_dataset = dataset[:split_index]
-# test_dataset = dataset[split_index:]
+train_dataset, val_dataset, test_dataset = random_split(
+    all_dataset, [train_size, val_size, test_size])
 
 
 train_loader = DataLoader(
     train_dataset, batch_size=config.batch_size, shuffle=True)
 
+val_loader = DataLoader(
+    val_dataset, batch_size=config.batch_size, shuffle=True)
+
+test_loader = DataLoader(
+    test_dataset, batch_size=config.batch_size, shuffle=True)
+
 
 def get_train_dataset(batch_size):
     return DataLoader(
-        train_dataset, batch_size=batch_size, shuffle=True)
-
+        all_dataset, batch_size=batch_size, shuffle=True)