From b03404ba48c71f6f734492c8598f66969d2c7495 Mon Sep 17 00:00:00 2001 From: hjdeheer Date: Tue, 16 Apr 2024 14:45:53 +0200 Subject: [PATCH] Change torch trainer to numpy and rename n_folds --- .../pipeline/model/training/torch_trainer.py | 41 ++++++++------- .../model/training/test_torch_trainer.py | 52 ++++++++++++++++--- 2 files changed, 68 insertions(+), 25 deletions(-) diff --git a/epochalyst/pipeline/model/training/torch_trainer.py b/epochalyst/pipeline/model/training/torch_trainer.py index a1236ca..3b942ff 100644 --- a/epochalyst/pipeline/model/training/torch_trainer.py +++ b/epochalyst/pipeline/model/training/torch_trainer.py @@ -36,6 +36,8 @@ class TorchTrainer(TrainingBlock): - `batch_size` (int): Batch size - `patience` (int): Patience for early stopping - `test_size` (float): Relative size of the test set + - `n_folds` (float): Number of folds for cross validation (0 for train full, + - `fold` (int): Fold number Methods: .. code-block:: python @@ -125,14 +127,14 @@ def log_to_terminal(self, message: str) -> None: test_size: Annotated[float, Interval(ge=0, le=1)] = 0.2 # Hashing purposes _fold: int = field(default=-1, init=False, repr=False, compare=False) - test_split_type: float = field(default=-1, init=True, repr=False, compare=False) + n_folds: float = field(default=-1, init=True, repr=False, compare=False) def __post_init__(self) -> None: """Post init method for the TorchTrainer class.""" - if self.test_split_type == -1: + if self.n_folds == -1: raise ValueError( - "Train_split_type needs to be set to either test_size or n_folds" + "Please specify the number of folds for cross validation or set n_folds to 0 for train full." ) self.save_model_to_disk = True @@ -171,7 +173,7 @@ def custom_train( x: npt.NDArray[np.float32], y: npt.NDArray[np.float32], **train_args: Any, - ) -> tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[npt.NDArray[np.float32], npt.NDArray[np.float32]]: """Train the model. :param x: The input to the system. @@ -218,7 +220,7 @@ def custom_train( ) self._load_model() # Return the predictions - return self.predict_on_loader(pred_dataloader), torch.tensor(y) + return self.predict_on_loader(pred_dataloader), y self.log_to_terminal(f"Training model: {self.model.__class__.__name__}") self.log_to_debug(f"Training model: {self.model.__class__.__name__}") @@ -253,11 +255,11 @@ def custom_train( if save_model: self._save_model() - return self.predict_on_loader(pred_dataloader), torch.tensor(y) + return self.predict_on_loader(pred_dataloader), y def custom_predict( self, x: npt.NDArray[np.float32], **pred_args: Any - ) -> torch.Tensor: + ) -> npt.NDArray[np.float32]: """Predict on the test data :param x: The input to the system. @@ -277,26 +279,27 @@ def custom_predict( pred_dataset, batch_size=curr_batch_size, shuffle=False ) - # Predict with a single model, test_split_type lower than 1 means a single test size, no CV - if self.test_split_type < 1 or pred_args.get("use_single_model", False): + # Predict with a single model, n_folds lower than 1 means a single test size, no CV + if self.n_folds < 1 or pred_args.get("use_single_model", False): self._load_model() return self.predict_on_loader(pred_dataloader) # Ensemble the fold models: predictions = [] - for i in range(int(self.test_split_type)): - self.log_to_terminal( - f"Predicting with model fold {i + 1}/{self.test_split_type}" - ) + for i in range(int(self.n_folds)): + self.log_to_terminal(f"Predicting with model fold {i + 1}/{self.n_folds}") self._fold = i # set the fold, which updates the hash self._load_model() # load the model for this fold predictions.append(self.predict_on_loader(pred_dataloader)) - test_predictions = torch.stack(predictions) + # Average the predictions using numpy + test_predictions = np.array(predictions) - return torch.mean(test_predictions, dim=0) + return np.mean(test_predictions, axis=0) - def predict_on_loader(self, loader: DataLoader[tuple[Tensor, ...]]) -> torch.Tensor: + def predict_on_loader( + self, loader: DataLoader[tuple[Tensor, ...]] + ) -> npt.NDArray[np.float32]: """Predict on the loader. :param loader: The loader to predict on. @@ -309,11 +312,11 @@ def predict_on_loader(self, loader: DataLoader[tuple[Tensor, ...]]) -> torch.Ten for data in tepoch: X_batch = data[0].to(self.device).float() - y_pred = self.model(X_batch).cpu() + y_pred = self.model(X_batch).cpu().numpy() predictions.extend(y_pred) self.log_to_terminal("Done predicting") - return torch.stack(predictions) + return np.array(predictions) def get_hash(self) -> str: """Get the hash of the block. @@ -322,7 +325,7 @@ def get_hash(self) -> str: :return: The hash of the block. """ - result = f"{self._hash}_{self.test_split_type}" + result = f"{self._hash}_{self.n_folds}" if self._fold != -1: result += f"_f{self._fold}" return result diff --git a/tests/pipeline/model/training/test_torch_trainer.py b/tests/pipeline/model/training/test_torch_trainer.py index 291bba4..2d28f74 100644 --- a/tests/pipeline/model/training/test_torch_trainer.py +++ b/tests/pipeline/model/training/test_torch_trainer.py @@ -16,7 +16,7 @@ class TestTorchTrainer: class ImplementedTorchTrainer(TorchTrainer): def __post_init__(self): - self.test_split_type = 1 + self.n_folds = 1 super().__post_init__() def log_to_terminal(self, message: str) -> None: @@ -28,7 +28,7 @@ def log_to_debug(self, message: str) -> None: @dataclass class FullyImplementedTorchTrainer(TorchTrainer): def __post_init__(self): - self.test_split_type = 1 + self.n_folds = 1 super().__post_init__() def log_to_terminal(self, message: str) -> None: @@ -48,7 +48,7 @@ def log_to_warning(self, message: str) -> None: def test_init_no_args(self): with pytest.raises(TypeError): - TorchTrainer(test_split_type=1) + TorchTrainer(n_folds=1) def test_init_none_args(self): with pytest.raises(TypeError): @@ -57,7 +57,7 @@ def test_init_none_args(self): criterion=None, optimizer=None, device=None, - test_split_type=1, + n_folds=1, ) def test_init_proper_args(self): @@ -66,7 +66,7 @@ def test_init_proper_args(self): model=self.simple_model, criterion=torch.nn.MSELoss(), optimizer=self.optimizer, - test_split_type=0, + n_folds=0, ) def test_init_proper_args_with_implemented(self): @@ -195,7 +195,7 @@ def test_train_full(self): criterion=torch.nn.MSELoss(), optimizer=self.optimizer, ) - tt.test_split_type = 0 + tt.n_folds = 0 tt.update_model_directory("tests/cache") x = torch.rand(10, 1) y = torch.rand(10) @@ -245,6 +245,46 @@ def test_predict(self): remove_cache_files() + def test_predict_3fold(self): + tt = self.FullyImplementedTorchTrainer( + model=self.simple_model, + criterion=torch.nn.MSELoss(), + optimizer=self.optimizer, + ) + remove_cache_files() + tt.n_folds = 3 + tt.update_model_directory("tests/cache") + x = torch.rand(10, 1) + y = torch.rand(10) + tt.train( + x, y, train_indices=[0, 1, 2, 3, 4, 5, 6, 7], test_indices=[8, 9], fold=0 + ) + tt.train( + x, y, train_indices=[0, 1, 2, 3, 4, 5, 6, 7], test_indices=[8, 9], fold=1 + ) + tt.train( + x, y, train_indices=[0, 1, 2, 3, 4, 5, 6, 7], test_indices=[8, 9], fold=2 + ) + tt.predict(x) + + remove_cache_files() + + def test_predict_train_full(self): + tt = self.FullyImplementedTorchTrainer( + model=self.simple_model, + criterion=torch.nn.MSELoss(), + optimizer=self.optimizer, + ) + remove_cache_files() + tt.n_folds = 0 + tt.update_model_directory("tests/cache") + x = torch.rand(10, 1) + y = torch.rand(10) + tt.train(x, y, train_indices=[0, 1, 2, 3, 4, 5, 6, 7], test_indices=[]) + tt.predict(x) + + remove_cache_files() + def test_predict_no_model_trained(self): tt = self.FullyImplementedTorchTrainer( model=self.simple_model,