Streamline CNN data handling

Author: seofernando25

hparam_search.py

@@ -6,19 +6,14 @@
 
 import numpy as np
 import torch
-from torch.utils.data import DataLoader
 from torch.utils.tensorboard import SummaryWriter
 
 from framework.data_utils import (
-    create_dataloaders,
     load_cifar10_data,
     prepare_data,
     split_train_val,
 )
-from framework.datasets import CIFAR10Dataset
-from framework.utils import get_device
 from models.base import get_model_by_name
-from models.cnn import TrainingConfig
 from search import RandomSearch
 
 RANDOM_SEED = 321
@@ -81,42 +76,14 @@ def evaluate_model(
         return model.evaluate(data["val_flat"], data["val_labels"])
 
     if model_key == "cnn":
-        # Architecture specific parameters
-        architecture = {k: params[k] for k in ("kernel_size", "stride")}
-        model.create_model(**architecture)
-
-        # Training specific parameters
-        batch_size = int(params["batch_size"])
-        config = TrainingConfig(
-            epochs=DEFAULT_EPOCHS,
-            learning_rate=float(params["learning_rate"]),
-            weight_decay=float(params["weight_decay"]),
-            optimizer=params["optimizer"],
-            patience=DEFAULT_PATIENCE,
-            batch_size=batch_size,
-        )
-        train_loader, val_loader = create_dataloaders(
+        model.create_model(**params)
+        model.train(
             data["train_images"],
             data["train_labels"],
             data["val_images"],
             data["val_labels"],
-            batch_size=batch_size,
-        )
-
-        device = get_device()
-        model.train(
-            train_loader, val_loader, config=config, device=device
-        )
-
-        eval_loader = DataLoader(
-            CIFAR10Dataset(data["val_images"], data["val_labels"]),
-            batch_size=batch_size,
-            shuffle=False,
-            num_workers=0,
-            pin_memory=torch.cuda.is_available(),
         )
-        eval_metrics = model.evaluate(eval_loader, device=device)
-        return eval_metrics
+        return model.evaluate(data["val_images"], data["val_labels"])
 
     raise ValueError(f"Unsupported model key: {model_key}")
 

models/cnn.py

@@ -2,14 +2,17 @@
 
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Dict, Optional
+from typing import Dict, List, Optional
 
+import numpy as np
 import torch
 import torch.nn as nn
 import torch.optim as optim
 from torch.utils.data import DataLoader
 from torch.utils.tensorboard import SummaryWriter
 
+from framework.data_utils import create_dataloaders
+from framework.datasets import CIFAR10Dataset
 from framework.training import Checkpoint, EarlyStopping, train_epoch, validate
 from framework.utils import count_parameters, get_device
 from .ParamSpace import ParamSpace
@@ -106,17 +109,47 @@ def create_model(self, **params) -> None:
 
     def train(
         self,
-        train_loader: DataLoader,
-        val_loader: DataLoader,
-        config: Optional[TrainingConfig] = None,
+        X_train: List[np.ndarray],
+        y_train: np.ndarray,
+        X_val: List[np.ndarray],
+        y_val: np.ndarray,
         device: Optional[torch.device] = None,
+        epochs: Optional[int] = None,
+        patience: Optional[int] = None,
+        min_delta: Optional[float] = None,
+        checkpoint_path: Optional[Path] = None,
+        grad_clip_norm: Optional[float] = None,
+        writer: Optional[SummaryWriter] = None,
+        num_workers: int = 2,
     ) -> Dict[str, float]:
         if self.network is None:
             raise RuntimeError("Train called before model is initialized")
         device = device or get_device()
         self.network = self.network.to(device)
 
-        config = config or TrainingConfig()
+        default_config = TrainingConfig()
+        config = TrainingConfig(
+            learning_rate=float(self.params.get("learning_rate", default_config.learning_rate)),
+            weight_decay=float(self.params.get("weight_decay", default_config.weight_decay)),
+            optimizer=self.params.get("optimizer", default_config.optimizer),
+            batch_size=int(self.params.get("batch_size", default_config.batch_size)),
+            # Infrastructure params: use provided values or defaults
+            epochs=epochs if epochs is not None else default_config.epochs,
+            patience=patience if patience is not None else default_config.patience,
+            min_delta=min_delta if min_delta is not None else default_config.min_delta,
+            checkpoint_path=checkpoint_path if checkpoint_path is not None else default_config.checkpoint_path,
+            grad_clip_norm=grad_clip_norm if grad_clip_norm is not None else default_config.grad_clip_norm,
+            writer=writer if writer is not None else default_config.writer,
+        )
+
+        train_loader, val_loader = create_dataloaders(
+            X_train,
+            y_train,
+            X_val,
+            y_val,
+            batch_size=config.batch_size,
+            num_workers=num_workers,
+        )
 
         optimizer = self._build_optimizer(self.network, config)
         scheduler = optim.lr_scheduler.OneCycleLR(
@@ -226,16 +259,29 @@ def predict(
 
     def evaluate(
         self,
-        data_loader: DataLoader,
+        X: List[np.ndarray],
+        y: np.ndarray,
         device: Optional[torch.device] = None,
         criterion: Optional[nn.Module] = None,
+        num_workers: int = 0,
     ) -> Dict[str, float]:
         if self.network is None:
             raise RuntimeError("Evaluate called before model is initialized")
         device = device or get_device()
         network = self.network.to(device)
         network.eval()
 
+        default_config = TrainingConfig()
+        batch_size = int(self.params.get("batch_size", default_config.batch_size))
+        dataset = CIFAR10Dataset(X, y)
+        data_loader = DataLoader(
+            dataset,
+            batch_size=batch_size,
+            shuffle=False,
+            num_workers=num_workers,
+            pin_memory=torch.cuda.is_available(),
+        )
+
         criterion = criterion or nn.CrossEntropyLoss()
 
         total_loss = 0.0

scripts/train_cnn.py

@@ -8,13 +8,12 @@
 
 sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
 from framework.data_utils import (
-    create_dataloaders,
     load_cifar10_data,
     prepare_data,
     split_train_val,
 )
 from framework.utils import get_device, test_pytorch_setup
-from models.cnn import CNNModel, TrainingConfig
+from models.cnn import CNNModel
 
 
 def parse_args():
@@ -114,34 +113,32 @@ def train_model(args, writer: SummaryWriter):
     )
     print(f"Train samples: {len(X_train)}, Val samples: {len(X_val)}")
 
-    train_loader, val_loader = create_dataloaders(
+    model = CNNModel(num_classes=num_classes)
+    # Pass hyperparameters via create_model (stored in model.params)
+    model.create_model(
+        learning_rate=args.lr,
+        weight_decay=args.weight_decay,
+        optimizer=args.optimizer,
+        batch_size=args.batch_size,
+    )
+
+    test_pytorch_setup()
+    # train() creates DataLoaders internally using batch_size from model.params
+    results = model.train(
         X_train,
         y_train,
         X_val,
         y_val,
-        batch_size=args.batch_size,
-        num_workers=args.num_workers,
-    )
-
-    model = CNNModel(num_classes=num_classes)
-    model.create_model()
-
-    config = TrainingConfig(
+        device=device,
         epochs=args.epochs,
-        learning_rate=args.lr,
-        weight_decay=args.weight_decay,
-        optimizer=args.optimizer,
         patience=args.patience,
         min_delta=args.min_delta,
         checkpoint_path=Path(args.checkpoint_path),
         grad_clip_norm=args.grad_clip,
-        batch_size=args.batch_size,
         writer=writer,
+        num_workers=args.num_workers,
     )
 
-    test_pytorch_setup()
-    results = model.train(train_loader, val_loader, config=config, device=device)
-
     print("\nTraining complete!")
     print(
         f"Best val acc: {results['best_val_acc']:.4f} ({results['best_val_acc'] * 100:.2f}%)"