change default epochs in train

Author: ArthurDeclercq

deeplc/_model_ops.py

@@ -21,18 +21,26 @@ def load_model(
 ) -> torch.nn.Module:
     """Load a model from a file or return a randomly initialized model if none is provided."""
     # If device is not specified, use the default device (GPU if available, else CPU)
-    selected_device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    selected_device = device or torch.device(
+        "cuda" if torch.cuda.is_available() else "cpu"
+    )
 
     # Load model from file if a path is provided
     if isinstance(model, str | Path):
-        loaded_model = torch.load(model, weights_only=False, map_location=selected_device)
+        loaded_model = torch.load(
+            model, weights_only=False, map_location=selected_device
+        )
     elif isinstance(model, torch.nn.Module):
         loaded_model = model
     elif model is None:
         # TODO: Implement randomly initialized model; requires model architecture definition
-        raise NotImplementedError("Loading randomly initialized model is not implemented yet.")
+        raise NotImplementedError(
+            "Loading randomly initialized model is not implemented yet."
+        )
     else:
-        raise TypeError(f"Expected a PyTorch Module or a file path, got {type(model)} instead.")
+        raise TypeError(
+            f"Expected a PyTorch Module or a file path, got {type(model)} instead."
+        )
 
     # Ensure the model is on the specified device
     loaded_model.to(selected_device)
@@ -47,7 +55,7 @@ def train(
     validation_split: float = 0.1,
     device: str = "cpu",
     learning_rate: float = 0.001,
-    epochs: int = 10,
+    epochs: int = 25,
     batch_size: int = 512,
     patience: int = 5,
     trainable_layers: str | None = None,
@@ -93,9 +101,15 @@ def train(
         logger.debug(f"Frozen all layers except those containing '{trainable_layers}'")
 
     # Parse dataset and split arguments; setup loaders
-    train_dataset, val_dataset = split_datasets(train_data, validation_data, validation_split)
-    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=0)
-    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=0)
+    train_dataset, val_dataset = split_datasets(
+        train_data, validation_data, validation_split
+    )
+    train_loader = DataLoader(
+        train_dataset, batch_size=batch_size, shuffle=True, num_workers=0
+    )
+    val_loader = DataLoader(
+        val_dataset, batch_size=batch_size, shuffle=False, num_workers=0
+    )
 
     optimizer = _get_optimizer(model, learning_rate)
     loss_fn = torch.nn.L1Loss()
@@ -137,7 +151,9 @@ def predict(
 ) -> torch.Tensor:
     """Predict using the model for the given dataset."""
     model = load_model(model, device)
-    data_loader = DataLoader(data, batch_size=batch_size, shuffle=False, num_workers=num_workers)
+    data_loader = DataLoader(
+        data, batch_size=batch_size, shuffle=False, num_workers=num_workers
+    )
     predictions = _predict_epoch(model, data_loader, device)
     return predictions.cpu().detach()
 
@@ -151,7 +167,9 @@ def evaluate(
 ) -> float:
     """Evaluate the model on the given dataset."""
     model = load_model(model, device)
-    data_loader = DataLoader(data, batch_size=batch_size, shuffle=False, num_workers=num_workers)
+    data_loader = DataLoader(
+        data, batch_size=batch_size, shuffle=False, num_workers=num_workers
+    )
     loss_fn = torch.nn.L1Loss()
     avg_loss = _validate_epoch(model, data_loader, loss_fn, device)
     return avg_loss
@@ -166,7 +184,9 @@ def _freeze_layers(model: torch.nn.Module, unfreeze_keyword: str) -> None:
         param.requires_grad = unfreeze_keyword in name
 
 
-def _get_optimizer(model: torch.nn.Module, learning_rate: float) -> torch.optim.Optimizer:
+def _get_optimizer(
+    model: torch.nn.Module, learning_rate: float
+) -> torch.optim.Optimizer:
     return torch.optim.Adam(
         filter(lambda p: p.requires_grad, model.parameters()),
         lr=learning_rate,