diff --git a/bluemath_tk/deeplearning/_base_model.py b/bluemath_tk/deeplearning/_base_model.py index c764f4e..7b406ee 100644 --- a/bluemath_tk/deeplearning/_base_model.py +++ b/bluemath_tk/deeplearning/_base_model.py @@ -1,6 +1,6 @@ import copy +import inspect from abc import abstractmethod -from typing import Dict, Optional, Union import numpy as np import torch @@ -8,6 +8,8 @@ from tqdm import tqdm from ..core.models import BlueMathModel +from .metrics import evaluate_reconstruction as evaluate_reconstruction_metric +from .metrics import reconstruction_error as reconstruction_error_metric class BaseDeepLearningModel(BlueMathModel): @@ -28,7 +30,7 @@ class BaseDeepLearningModel(BlueMathModel): """ @abstractmethod - def __init__(self, device: Optional[Union[str, torch.device]] = None, **kwargs): + def __init__(self, device: str | torch.device | None = None, **kwargs): """ Initialize the base deep learning model. @@ -51,8 +53,9 @@ def __init__(self, device: Optional[Union[str, torch.device]] = None, **kwargs): else: self.device = device - self.model: Optional[nn.Module] = None + self.model: nn.Module | None = None self.is_fitted = False + self._build_input_shape: tuple | None = None self._exclude_attributes = [ "model", @@ -71,20 +74,176 @@ def _build_model(self, *args, **kwargs) -> nn.Module: pass + + def _get_reconstruction_target(self, X: np.ndarray) -> np.ndarray: + """Return the default reconstruction target for ``X``.""" + return X + + @staticmethod + def _batch_slices( + n_samples: int, + batch_size: int, + avoid_singleton: bool = False, + ) -> list[tuple[int, int]]: + """Return batch slices, optionally avoiding a final singleton batch.""" + if avoid_singleton and batch_size == 1: + raise ValueError( + "batch_size=1 is incompatible with this model's BatchNorm1d " + "layers. Use batch_size>=2." + ) + + slices = [ + (start, min(start + batch_size, n_samples)) + for start in range(0, n_samples, batch_size) + ] + + if ( + avoid_singleton + and len(slices) > 1 + and slices[-1][1] - slices[-1][0] == 1 + ): + previous_start, previous_stop = slices[-2] + final_stop = slices[-1][1] + previous_size = previous_stop - previous_start + + if previous_size > 2: + slices[-2] = (previous_start, previous_stop - 1) + slices[-1] = (previous_stop - 1, final_stop) + else: + slices[-2] = (previous_start, final_stop) + slices.pop() + + return slices + + def _requires_non_singleton_training_batches(self) -> bool: + """Return whether BatchNorm1d requires at least two training samples.""" + if self.model is None: + return False + return any( + isinstance(module, nn.BatchNorm1d) + for module in self.model.modules() + ) + + def _validate_or_set_build_input_shape(self, input_shape: tuple) -> None: + """Store the build shape or reject incompatible repeated fitting.""" + input_shape = tuple(input_shape) + if self._build_input_shape is None: + self._build_input_shape = input_shape + return + + if tuple(self._build_input_shape[1:]) != tuple(input_shape[1:]): + raise ValueError( + "The fitted model input shape is incompatible with the new data. " + f"Expected per-sample shape {self._build_input_shape[1:]}, " + f"got {input_shape[1:]}." + ) + + @staticmethod + def _validate_inference_inputs( + X: np.ndarray, + batch_size: int, + name: str = "X", + ) -> None: + """Validate common prediction, encoding, and decoding inputs.""" + if not isinstance(X, np.ndarray): + raise TypeError(f"{name} must be a NumPy array.") + if X.ndim < 1: + raise ValueError(f"{name} must have at least one dimension.") + if len(X) == 0: + raise ValueError(f"{name} must contain at least one sample.") + if batch_size < 1: + raise ValueError("batch_size must be at least 1.") + + @staticmethod + def _validate_fit_inputs( + X: np.ndarray, + y: np.ndarray, + validation_split: float, + batch_size: int, + epochs: int, + patience: int, + ) -> None: + """Validate common training inputs before building the model.""" + if not isinstance(X, np.ndarray): + raise TypeError("X must be a NumPy array.") + if not isinstance(y, np.ndarray): + raise TypeError("y must be a NumPy array.") + if X.ndim < 2: + raise ValueError( + "X must include a leading sample dimension. " + "For tabular data use shape (n_samples, n_features)." + ) + if len(X) != len(y): + raise ValueError( + f"X and y must contain the same number of samples; " + f"got {len(X)} and {len(y)}." + ) + if not 0.0 < validation_split < 1.0: + raise ValueError("validation_split must be strictly between 0 and 1.") + if batch_size < 1: + raise ValueError("batch_size must be at least 1.") + if epochs < 1: + raise ValueError("epochs must be at least 1.") + if patience < 1: + raise ValueError("patience must be at least 1.") + + split = int((1 - validation_split) * len(X)) + if split < 2: + raise ValueError( + "The training split must contain at least two samples. " + "Increase the dataset size or reduce validation_split." + ) + if len(X) - split < 1: + raise ValueError( + "The validation split must contain at least one sample." + ) + + def _get_init_config(self) -> dict: + """Collect constructor parameters needed to recreate this model.""" + config = {} + missing = [] + signature = inspect.signature(self.__class__.__init__) + for name, parameter in signature.parameters.items(): + if name in {"self", "device"}: + continue + if parameter.kind is inspect.Parameter.VAR_KEYWORD: + continue + if hasattr(self, name): + config[name] = getattr(self, name) + else: + missing.append(name) + + if missing: + raise ValueError( + f"{self.__class__.__name__} cannot create a self-describing " + f"checkpoint because these constructor parameters are not " + f"stored as attributes: {missing}. Override _get_init_config()." + ) + return config + + @staticmethod + def _require_scalar_loss(loss: torch.Tensor) -> None: + """Raise when a criterion returns a non-scalar training loss.""" + if loss.ndim != 0: + raise ValueError( + "The training criterion must return a scalar loss. " + "Use reduction='mean' or reduction='sum'." + ) + def fit( self, X: np.ndarray, - y: Optional[np.ndarray] = None, + y: np.ndarray | None = None, validation_split: float = 0.2, epochs: int = 500, batch_size: int = 64, learning_rate: float = 1e-3, - optimizer: Optional[torch.optim.Optimizer] = None, - criterion: Optional[nn.Module] = None, + optimizer: torch.optim.Optimizer | None = None, + criterion: nn.Module | None = None, patience: int = 20, verbose: int = 1, **kwargs, - ) -> Dict[str, list]: + ) -> dict[str, list]: """ Fit the model. @@ -115,14 +274,31 @@ def fit( Returns ------- - Dict[str, list] + dict[str, list] Training history with 'train_loss' and 'val_loss' keys. """ + if not isinstance(X, np.ndarray): + raise TypeError("X must be a NumPy array.") + if y is None: + y = self._get_reconstruction_target(X) + + self._validate_fit_inputs( + X, + y, + validation_split, + batch_size, + epochs, + patience, + ) + self._validate_or_set_build_input_shape(tuple(X.shape)) + if self.model is None: self.model = self._build_model(X.shape, **kwargs) self.model = self.model.to(self.device) + avoid_singleton = self._requires_non_singleton_training_batches() + if optimizer is None: optimizer = torch.optim.Adam(self.model.parameters(), lr=learning_rate) @@ -137,11 +313,7 @@ def fit( train_idx, val_idx = idx[:split], idx[split:] Xtr, Xval = X[train_idx], X[val_idx] - if y is None: - # Autoencoder case - ytr, yval = Xtr, Xval - else: - ytr, yval = y[train_idx], y[val_idx] + ytr, yval = y[train_idx], y[val_idx] # Convert to tensors Xtr_tensor = torch.FloatTensor(Xtr).to(self.device) @@ -166,15 +338,21 @@ def fit( # Training self.model.train() train_loss = 0.0 - n_batches = (len(Xtr) + batch_size - 1) // batch_size + train_slices = self._batch_slices( + len(Xtr), + batch_size, + avoid_singleton=avoid_singleton, + ) + n_batches = len(train_slices) - for i in range(0, len(Xtr), batch_size): - batch_X = Xtr_tensor[i : i + batch_size] - batch_y = ytr_tensor[i : i + batch_size] + for start, stop in train_slices: + batch_X = Xtr_tensor[start:stop] + batch_y = ytr_tensor[start:stop] optimizer.zero_grad() output = self.model(batch_X) loss = criterion(output, batch_y) + self._require_scalar_loss(loss) loss.backward() optimizer.step() @@ -187,13 +365,19 @@ def fit( self.model.eval() val_loss = 0.0 with torch.no_grad(): - n_val_batches = (len(Xval) + batch_size - 1) // batch_size - for i in range(0, len(Xval), batch_size): - batch_X = Xval_tensor[i : i + batch_size] - batch_y = yval_tensor[i : i + batch_size] + val_slices = self._batch_slices( + len(Xval), + batch_size, + avoid_singleton=False, + ) + n_val_batches = len(val_slices) + for start, stop in val_slices: + batch_X = Xval_tensor[start:stop] + batch_y = yval_tensor[start:stop] output = self.model(batch_X) loss = criterion(output, batch_y) + self._require_scalar_loss(loss) val_loss += loss.item() val_loss /= n_val_batches @@ -216,11 +400,15 @@ def fit( # Update progress bar with current losses if pbar is not None: pbar.set_postfix_str( - f"Train Loss: {train_loss:.6f}, Val Loss: {val_loss:.6f}, Patience: {patience_counter}/{patience}" + f"Train Loss: {train_loss:.6f}, " + f"Val Loss: {val_loss:.6f}, " + f"Patience: {patience_counter}/{patience}" ) elif verbose > 0 and (epoch + 1) % max(1, epochs // 10) == 0: self.logger.info( - f"Epoch {epoch + 1}/{epochs} - Train Loss: {train_loss:.6f}, Val Loss: {val_loss:.6f}" + f"Epoch {epoch + 1}/{epochs} - " + f"Train Loss: {train_loss:.6f}, " + f"Val Loss: {val_loss:.6f}" ) # Restore best model @@ -259,6 +447,7 @@ def predict( if not self.is_fitted or self.model is None: raise ValueError("Model must be fitted before prediction.") + self._validate_inference_inputs(X, batch_size) self.model.eval() X_tensor = torch.FloatTensor(X).to(self.device) @@ -308,6 +497,7 @@ def encode( if not self.is_fitted or self.model is None: raise ValueError("Model must be fitted before encoding.") + self._validate_inference_inputs(X, batch_size) # Check if model has encode_forward method if not hasattr(self.model, "encode_forward"): @@ -336,48 +526,198 @@ def encode( return np.concatenate(encodings, axis=0) - def save_pytorch_model(self, model_path: str, **kwargs): - """ - Save the PyTorch model to a file. - Parameters - ---------- - model_path : str - Path to the file where the model will be saved. - **kwargs - Additional arguments for torch.save. - """ + def decode( + self, + Z: np.ndarray, + batch_size: int = 64, + verbose: int = 1, + ) -> np.ndarray: + """Decode latent vectors into the reconstruction space.""" + if not self.is_fitted or self.model is None: + raise ValueError("Model must be fitted before decoding.") + if not hasattr(self.model, "decode_forward"): + raise ValueError( + f"Model {self.__class__.__name__} does not support decoding." + ) + + Z = np.asarray(Z) + if Z.ndim == 1: + Z = Z[None, :] + self._validate_inference_inputs(Z, batch_size, name="Z") + + self.model.eval() + Z_tensor = torch.as_tensor(Z, dtype=torch.float32, device=self.device) + outputs = [] + batch_range = range(0, len(Z), batch_size) + n_batches = (len(Z) + batch_size - 1) // batch_size + + if verbose > 0 and n_batches > 1: + batch_range = tqdm( + batch_range, + desc="Decoding", + unit="batch", + total=n_batches, + ) + + with torch.no_grad(): + for start in batch_range: + output = self.model.decode_forward( + Z_tensor[start : start + batch_size] + ) + outputs.append(output.cpu().numpy()) + + return np.concatenate(outputs, axis=0) + def reconstruction_error( + self, + X: np.ndarray, + y: np.ndarray | None = None, + metric: str = "mse", + reduction: str = "sample", + batch_size: int = 64, + verbose: int = 0, + eps: float = 0.0, + ): + """Compute reconstruction error with the shared metrics module.""" + self._validate_inference_inputs(X, batch_size) + target = self._get_reconstruction_target(X) if y is None else y + prediction = self.predict(X, batch_size=batch_size, verbose=verbose) + return reconstruction_error_metric( + target, + prediction, + metric=metric, + reduction=reduction, + eps=eps, + ) + + def evaluate_reconstruction( + self, + X: np.ndarray, + y: np.ndarray | None = None, + metric: str = "mse", + batch_size: int = 64, + verbose: int = 0, + eps: float = 0.0, + ) -> dict[str, float]: + """Return summary statistics for reconstruction error.""" + self._validate_inference_inputs(X, batch_size) + target = self._get_reconstruction_target(X) if y is None else y + prediction = self.predict(X, batch_size=batch_size, verbose=verbose) + return evaluate_reconstruction_metric( + target, + prediction, + metric=metric, + eps=eps, + ) + + def evaluate(self, X: np.ndarray, **kwargs) -> dict[str, float]: + """Alias for :meth:`evaluate_reconstruction`.""" + return self.evaluate_reconstruction(X, **kwargs) + + def save_pytorch_model(self, model_path: str, **kwargs): + """Save weights and metadata required to rebuild the model.""" if self.model is None: raise ValueError("PyTorch model must be built before saving.") + if self._build_input_shape is None: + raise ValueError( + "The model input shape is unknown. Fit or load the model before saving." + ) torch.save( { + "checkpoint_version": 2, "model_state_dict": self.model.state_dict(), "is_fitted": self.is_fitted, "model_class": self.__class__.__name__, + "init_config": self._get_init_config(), + "build_input_shape": self._build_input_shape, }, model_path, **kwargs, ) self.logger.info(f"PyTorch model saved to {model_path}") - def load_pytorch_model(self, model_path: str, **kwargs): - """ - Load a PyTorch model from a file. + def load_pytorch_model( + self, + model_path: str, + map_location=None, + **kwargs, + ): + """Load a checkpoint into this instance.""" + if map_location is None: + map_location = self.device - Parameters - ---------- - model_path : str - Path to the file where the model is saved. - **kwargs - Additional arguments for torch.load. - """ + checkpoint = torch.load( + model_path, + map_location=map_location, + **kwargs, + ) + checkpoint_class = checkpoint.get("model_class") + if checkpoint_class and checkpoint_class != self.__class__.__name__: + raise ValueError( + f"Checkpoint contains {checkpoint_class}, " + f"not {self.__class__.__name__}." + ) if self.model is None: - raise ValueError("PyTorch model must be built before loading.") + build_input_shape = checkpoint.get("build_input_shape") + if build_input_shape is None: + raise ValueError( + "This legacy checkpoint does not include build_input_shape. " + "Build the model manually before loading it." + ) + + init_config = checkpoint.get("init_config", {}) + for name, value in init_config.items(): + if hasattr(self, name): + setattr(self, name, value) + + self._build_input_shape = tuple(build_input_shape) + self.model = self._build_model(self._build_input_shape) + self.model = self.model.to(self.device) - checkpoint = torch.load(model_path, **kwargs) self.model.load_state_dict(checkpoint["model_state_dict"]) self.is_fitted = checkpoint.get("is_fitted", False) + if self._build_input_shape is None: + shape = checkpoint.get("build_input_shape") + if shape is not None: + self._build_input_shape = tuple(shape) self.logger.info(f"PyTorch model loaded from {model_path}") + return self + + @classmethod + def from_pytorch_model( + cls, + model_path: str, + device: str | torch.device | None = "cpu", + map_location="cpu", + **kwargs, + ): + """Create a fitted model from a self-describing checkpoint.""" + checkpoint = torch.load( + model_path, + map_location=map_location, + **kwargs, + ) + checkpoint_class = checkpoint.get("model_class") + if checkpoint_class and checkpoint_class != cls.__name__: + raise ValueError( + f"Checkpoint contains {checkpoint_class}, not {cls.__name__}." + ) + + init_config = checkpoint.get("init_config") + build_input_shape = checkpoint.get("build_input_shape") + if init_config is None or build_input_shape is None: + raise ValueError( + "Automatic loading requires a version-2 checkpoint containing " + "init_config and build_input_shape." + ) + + instance = cls(device=device, **dict(init_config)) + instance._build_input_shape = tuple(build_input_shape) + instance.model = instance._build_model(instance._build_input_shape) + instance.model = instance.model.to(instance.device) + instance.model.load_state_dict(checkpoint["model_state_dict"]) + instance.is_fitted = checkpoint.get("is_fitted", False) + return instance diff --git a/bluemath_tk/deeplearning/autoencoders.py b/bluemath_tk/deeplearning/autoencoders.py index 6583db8..1fe21e2 100644 --- a/bluemath_tk/deeplearning/autoencoders.py +++ b/bluemath_tk/deeplearning/autoencoders.py @@ -50,10 +50,10 @@ class StandardAutoencoder(BaseDeepLearningModel): Input Shape ----------- X : np.ndarray - Input data with shape (n_samples, n_features) or (n_samples,). - - For 2D arrays: (n_samples, n_features) - each row is a sample - - For 1D arrays: (n_features,) - single sample (will be reshaped) - The model automatically flattens multi-dimensional inputs. + Input data with a leading sample dimension. + For tabular data use (n_samples, n_features). Higher-dimensional + per-sample inputs are flattened internally and reconstructed to + their original sample shape. Examples -------- @@ -146,6 +146,12 @@ def encode_forward(self, x): x = x.unsqueeze(0) return self.encoder(x) + + def decode_forward(self, z): + """Decode latent vectors to the original sample shape.""" + x_recon = self.decoder(z) + return x_recon.view(x_recon.size(0), *self.sample_shape) + return StandardAutoencoderModel(n_features, self.hidden_dims, self.k, sample_shape) @@ -160,10 +166,10 @@ class OrthogonalAutoencoder(BaseDeepLearningModel): Input Shape ----------- X : np.ndarray - Input data with shape (n_samples, n_features) or (n_samples,). - - For 2D arrays: (n_samples, n_features) - each row is a sample - - For 1D arrays: (n_features,) - single sample (will be reshaped) - The model automatically flattens multi-dimensional inputs. + Input data with a leading sample dimension. + For tabular data use (n_samples, n_features). Higher-dimensional + per-sample inputs are flattened internally and reconstructed to + their original sample shape. Examples -------- @@ -296,6 +302,12 @@ def get_regularization_losses(self): decorr_loss = getattr(self.latent_decorr, "_loss", None) return ortho_loss, decorr_loss + + def decode_forward(self, z): + """Decode latent vectors to the original sample shape.""" + x_recon = self.decoder(z) + return x_recon.view(x_recon.size(0), *self.sample_shape) + return OrthogonalAutoencoderModel( n_features, self.hidden_dims, @@ -325,10 +337,26 @@ def fit( This method overrides the base fit() to properly add orthogonality and decorrelation regularization losses during training. """ + if not isinstance(X, np.ndarray): + raise TypeError("X must be a NumPy array.") + if y is None: + y = self._get_reconstruction_target(X) + self._validate_fit_inputs( + X, + y, + validation_split, + batch_size, + epochs, + patience, + ) + self._validate_or_set_build_input_shape(tuple(X.shape)) + if self.model is None: self.model = self._build_model(X.shape, **kwargs) self.model = self.model.to(self.device) + avoid_singleton = self._requires_non_singleton_training_batches() + if optimizer is None: optimizer = torch.optim.Adam(self.model.parameters(), lr=learning_rate) @@ -372,11 +400,15 @@ def fit( # Training self.model.train() train_loss = 0.0 - n_batches = (len(Xtr) + batch_size - 1) // batch_size - - for i in range(0, len(Xtr), batch_size): - batch_X = Xtr_tensor[i : i + batch_size] - batch_y = ytr_tensor[i : i + batch_size] + train_slices = self._batch_slices( + len(Xtr), + batch_size, + avoid_singleton=avoid_singleton, + ) + n_batches = len(train_slices) + for start, stop in train_slices: + batch_X = Xtr_tensor[start:stop] + batch_y = ytr_tensor[start:stop] optimizer.zero_grad() output = self.model(batch_X) @@ -389,6 +421,7 @@ def fit( if decorr_loss is not None: loss = loss + decorr_loss + self._require_scalar_loss(loss) loss.backward() optimizer.step() @@ -401,10 +434,15 @@ def fit( self.model.eval() val_loss = 0.0 with torch.no_grad(): - n_val_batches = (len(Xval) + batch_size - 1) // batch_size - for i in range(0, len(Xval), batch_size): - batch_X = Xval_tensor[i : i + batch_size] - batch_y = yval_tensor[i : i + batch_size] + val_slices = self._batch_slices( + len(Xval), + batch_size, + avoid_singleton=False, + ) + n_val_batches = len(val_slices) + for start, stop in val_slices: + batch_X = Xval_tensor[start:stop] + batch_y = yval_tensor[start:stop] output = self.model(batch_X) loss = criterion(output, batch_y) @@ -416,6 +454,7 @@ def fit( if decorr_loss is not None: loss = loss + decorr_loss + self._require_scalar_loss(loss) val_loss += loss.item() val_loss /= n_val_batches @@ -559,6 +598,18 @@ def encode_forward(self, x): z = self.latent(x[:, -1, :]) # Take last timestep return z + + def decode_forward(self, z): + """Decode latent vectors to full temporal sequences.""" + z_expanded = ( + self.latent_to_seq(z) + .unsqueeze(1) + .repeat(1, self.seq_len, 1) + ) + x, _ = self.lstm3(z_expanded) + x, _ = self.lstm4(x) + return x + return LSTMAutoencoderModel(seq_len, n_features, self.hidden, self.k) @@ -740,6 +791,23 @@ def encode_forward(self, x): return z + + def decode_forward(self, z): + """Decode latent vectors to channels-first spatial grids.""" + batch_size = z.size(0) + x = F.relu(self.fc3(z)) + x = F.relu(self.fc4(x)) + x = x.view( + batch_size, + 64, + (self.H + self.pad_h) // 4, + (self.W + self.pad_w) // 4, + ) + x = self.decoder(x) + if self.pad_h > 0 or self.pad_w > 0: + x = x[:, :, : self.H, : self.W] + return x + return CNNAutoencoderModel(H, W, C, self.k, pad_h, pad_w) @@ -993,6 +1061,24 @@ def encode_forward(self, x): return z_k + + def decode_forward(self, z): + """Decode latent vectors to channels-first spatial grids.""" + batch_size = z.size(0) + dec_seed = F.relu(self.dec_seed(z)) + dec_tokens = dec_seed.view(batch_size, N, self.d_model) + dec_tokens = self.dec_pos_embed(dec_tokens) + y = dec_tokens + for block in self.decoder_blocks: + y = block(y) + patch_tokens = self.patch_reconstruct(y) + reconstruction = self.unpatchify(patch_tokens) + if self.pad_h > 0 or self.pad_w > 0: + reconstruction = reconstruction[ + :, :, : self.H, : self.W + ] + return reconstruction + return ViTAutoencoderModel( H, W, @@ -1095,6 +1181,16 @@ def fit( **kwargs, ) + + def _get_reconstruction_target(self, X: np.ndarray) -> np.ndarray: + """Use the last input frame as the default reconstruction target.""" + if X.ndim != 5: + raise ValueError( + "ConvLSTMAutoencoder expects 5D input " + "(n_samples, seq_len, C, H, W)." + ) + return X[:, -1] + def _build_model(self, input_shape: Tuple, **kwargs) -> nn.Module: """Build the ConvLSTM autoencoder model.""" # Parse input shape: (n_samples, seq_len, C, H, W) - channels-first format @@ -1270,6 +1366,25 @@ def encode_forward(self, x): return z + + def decode_forward(self, z): + """Decode latent vectors to the reconstructed final frame.""" + batch_size = z.size(0) + x = F.relu(self.fc_dec(z)) + x = x.view( + batch_size, + 64, + (self.H + self.pad_h) // 4, + (self.W + self.pad_w) // 4, + ) + x = self.upsample1(x) + x = F.relu(self.deconv1(x)) + x = self.upsample2(x) + x = self.deconv2(x) + if self.pad_h > 0 or self.pad_w > 0: + x = x[:, :, : self.H, : self.W] + return x + return ConvLSTMAutoencoderModel(seq_len, H, W, C, self.k, pad_h, pad_w) @@ -1376,6 +1491,16 @@ def fit( **kwargs, ) + + def _get_reconstruction_target(self, X: np.ndarray) -> np.ndarray: + """Use the last input frame as the default reconstruction target.""" + if X.ndim != 5: + raise ValueError( + "HybridConvLSTMTransformerAutoencoder expects 5D input " + "(n_samples, seq_len, C, H, W)." + ) + return X[:, -1] + def _build_model(self, input_shape: Tuple, **kwargs) -> nn.Module: """Build the hybrid autoencoder model.""" # Parse input shape: (n_samples, seq_len, C, H, W) - channels-first format @@ -1635,6 +1760,22 @@ def encode_forward(self, x): return z + + def decode_forward(self, z): + """Decode latent vectors to the reconstructed final frame.""" + batch_size = z.size(0) + x = F.relu(self.fc_dec(z)) + height = (self.H + self.pad_h) // 4 + width = (self.W + self.pad_w) // 4 + x = x.view(batch_size, 64, height, width) + x = self.upsample1(x) + x = F.relu(self.deconv1(x)) + x = self.upsample2(x) + x = self.deconv2(x) + if self.pad_h > 0 or self.pad_w > 0: + x = x[:, :, : self.H, : self.W] + return x + return HybridAutoencoderModel( seq_len, H, diff --git a/bluemath_tk/deeplearning/layers.py b/bluemath_tk/deeplearning/layers.py index 9d63de4..ace9740 100644 --- a/bluemath_tk/deeplearning/layers.py +++ b/bluemath_tk/deeplearning/layers.py @@ -350,14 +350,22 @@ def forward(self, z): """ # z: (batch, k) zc = z - z.mean(dim=0, keepdim=True) # center - B = zc.size(0) - cov = torch.matmul(zc.t(), zc) / (B - 1.0) # (k, k) - diag_mask = torch.eye(cov.size(0), device=cov.device, dtype=cov.dtype) - offdiag = cov * (1 - diag_mask) # zero diag - loss = self.strength * torch.sum(offdiag**2) + batch_size = zc.size(0) + + if batch_size < 2: + loss = z.new_zeros(()) + else: + cov = torch.matmul(zc.t(), zc) / (batch_size - 1.0) # (k, k) + diag_mask = torch.eye( + cov.size(0), + device=cov.device, + dtype=cov.dtype, + ) + offdiag = cov * (1 - diag_mask) + loss = self.strength * torch.sum(offdiag**2) # Add loss to computation graph - z = z + 0 * loss # Trick to add loss to graph without changing z + z = z + 0 * loss # Store current loss for retrieval during training self._loss = loss diff --git a/tests/deeplearning/test_autoencoder_api.py b/tests/deeplearning/test_autoencoder_api.py new file mode 100644 index 0000000..b1a72a6 --- /dev/null +++ b/tests/deeplearning/test_autoencoder_api.py @@ -0,0 +1,340 @@ +"""Tests for the public autoencoder API and checkpoint behaviour.""" + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") + +from bluemath_tk.deeplearning.autoencoders import ( # noqa: E402 + CNNAutoencoder, + ConvLSTMAutoencoder, + HybridConvLSTMTransformerAutoencoder, + LSTMAutoencoder, + OrthogonalAutoencoder, + StandardAutoencoder, + VisionTransformerAutoencoder, +) + + +@pytest.fixture(autouse=True) +def _set_reproducible_seed(): + np.random.seed(321) + torch.manual_seed(321) + torch.set_num_threads(1) + + +def _fit_kwargs(): + return { + "epochs": 1, + "batch_size": 4, + "validation_split": 0.25, + "patience": 2, + "verbose": 0, + "learning_rate": 1e-3, + } + + +@pytest.mark.parametrize( + ("factory", "shape", "decoded_shape"), + [ + ( + lambda: StandardAutoencoder(k=3, hidden_dims=[8], device="cpu"), + (16, 10), + (16, 10), + ), + ( + lambda: OrthogonalAutoencoder( + k=3, + hidden_dims=[8], + lambda_W=1e-4, + lambda_Z=1e-4, + device="cpu", + ), + (16, 10), + (16, 10), + ), + ( + lambda: LSTMAutoencoder(k=4, hidden=(8, 6), device="cpu"), + (16, 5, 3), + (16, 5, 3), + ), + ( + lambda: CNNAutoencoder(k=4, device="cpu"), + (16, 1, 8, 8), + (16, 1, 8, 8), + ), + ( + lambda: VisionTransformerAutoencoder( + k=4, + patch_size=4, + d_model=8, + depth_enc=1, + depth_dec=1, + heads=2, + device="cpu", + ), + (16, 1, 8, 8), + (16, 1, 8, 8), + ), + ( + lambda: ConvLSTMAutoencoder(k=4, device="cpu"), + (16, 3, 1, 8, 8), + (16, 1, 8, 8), + ), + ( + lambda: HybridConvLSTMTransformerAutoencoder( + k=4, + d_model=8, + n_heads=2, + n_layers=1, + efficient_attention="linear", + device="cpu", + ), + (16, 3, 1, 8, 8), + (16, 1, 8, 8), + ), + ], +) +def test_encode_decode_round_trip_shapes(factory, shape, decoded_shape): + """All autoencoders should decode latent vectors to the documented shape.""" + X = np.random.randn(*shape).astype("float32") + model = factory() + model.fit(X, **_fit_kwargs()) + + prediction = model.predict(X, batch_size=4, verbose=0) + Z = model.encode(X, batch_size=4, verbose=0) + decoded = model.decode(Z, batch_size=4, verbose=0) + + assert decoded.shape == decoded_shape + assert np.isfinite(decoded).all() + assert np.allclose(decoded, prediction, rtol=1e-5, atol=1e-6) + + +def test_standard_autoencoder_handles_singleton_remainder_batch(): + """StandardAutoencoder should avoid a singleton BatchNorm training batch.""" + X = np.random.randn(18, 10).astype("float32") + model = StandardAutoencoder(k=3, hidden_dims=[8], device="cpu") + history = model.fit(X, **_fit_kwargs()) + assert np.isfinite(history["train_loss"]).all() + + +def test_orthogonal_autoencoder_handles_singleton_remainder_batch(): + """OrthogonalAutoencoder should avoid singleton training batches.""" + X = np.random.randn(18, 10).astype("float32") + model = OrthogonalAutoencoder( + k=3, + hidden_dims=[8], + lambda_W=1e-4, + lambda_Z=1e-4, + device="cpu", + ) + history = model.fit(X, **_fit_kwargs()) + assert np.isfinite(history["train_loss"]).all() + + +def test_batch_size_one_is_rejected_for_batchnorm1d_autoencoders(): + """BatchNorm1d autoencoders should reject singleton training batches.""" + X = np.random.randn(8, 10).astype("float32") + model = StandardAutoencoder(k=3, hidden_dims=[8], device="cpu") + + with pytest.raises(ValueError, match="batch_size=1"): + model.fit( + X, + validation_split=0.25, + epochs=1, + batch_size=1, + patience=2, + verbose=0, + ) + + +def test_batch_size_one_remains_available_without_batchnorm1d(): + """Models without BatchNorm1d should retain explicit batch_size=1.""" + X = np.random.randn(6, 3, 2).astype("float32") + model = LSTMAutoencoder(k=2, hidden=(4, 3), device="cpu") + + history = model.fit( + X, + validation_split=0.33, + epochs=1, + batch_size=1, + patience=2, + verbose=0, + ) + + assert np.isfinite(history["train_loss"]).all() + + +def test_orthogonal_autoencoder_handles_singleton_validation_batch(): + """Latent decorrelation should be neutral for a one-sample validation batch.""" + X = np.random.randn(5, 6).astype("float32") + model = OrthogonalAutoencoder( + k=2, + hidden_dims=[4], + lambda_W=1e-4, + lambda_Z=1e-4, + device="cpu", + ) + + history = model.fit( + X, + validation_split=0.2, + epochs=1, + batch_size=2, + patience=2, + verbose=0, + ) + + assert np.isfinite(history["val_loss"]).all() + + +def test_repeated_fit_rejects_incompatible_sample_shape(): + """A built model should not silently accept a different sample shape.""" + X = np.random.randn(16, 10).astype("float32") + model = StandardAutoencoder(k=3, hidden_dims=[8], device="cpu") + model.fit(X, **_fit_kwargs()) + + X_incompatible = np.random.randn(16, 12).astype("float32") + with pytest.raises(ValueError, match="incompatible"): + model.fit(X_incompatible, **_fit_kwargs()) + + +def test_inference_methods_validate_batch_size(): + """Prediction, encoding, and decoding should reject invalid batch sizes.""" + X = np.random.randn(16, 10).astype("float32") + model = StandardAutoencoder(k=3, hidden_dims=[8], device="cpu") + model.fit(X, **_fit_kwargs()) + Z = model.encode(X, batch_size=4, verbose=0) + + with pytest.raises(ValueError, match="batch_size"): + model.predict(X, batch_size=0, verbose=0) + with pytest.raises(ValueError, match="batch_size"): + model.encode(X, batch_size=0, verbose=0) + with pytest.raises(ValueError, match="batch_size"): + model.decode(Z, batch_size=0, verbose=0) + + +def test_fit_validates_sample_dimension_and_validation_split(): + """Fit should reject ambiguous inputs and invalid validation splits.""" + model = StandardAutoencoder(k=2, hidden_dims=[4], device="cpu") + + with pytest.raises(ValueError, match="leading sample dimension"): + model.fit(np.ones(10, dtype="float32"), **_fit_kwargs()) + + X = np.ones((8, 3), dtype="float32") + with pytest.raises(ValueError, match="validation_split"): + model.fit( + X, + validation_split=0.0, + epochs=1, + batch_size=4, + patience=2, + verbose=0, + ) + + +def test_decode_accepts_one_latent_vector(): + """A one-dimensional latent vector should decode as one sample.""" + X = np.random.randn(16, 10).astype("float32") + model = StandardAutoencoder(k=3, hidden_dims=[8], device="cpu") + model.fit(X, **_fit_kwargs()) + Z = model.encode(X, batch_size=4, verbose=0) + + decoded = model.decode(Z[0], verbose=0) + + assert decoded.shape == (1, 10) + + +def test_standard_autoencoder_reconstruction_convenience_methods(): + """Convenience metrics should match direct reconstruction calculations.""" + X = np.random.randn(16, 10).astype("float32") + model = StandardAutoencoder(k=3, hidden_dims=[8], device="cpu") + model.fit(X, **_fit_kwargs()) + + prediction = model.predict(X, batch_size=4, verbose=0) + errors = model.reconstruction_error( + X, + metric="mse", + reduction="sample", + batch_size=4, + ) + summary = model.evaluate_reconstruction( + X, + metric="rmse", + batch_size=4, + ) + + manual = np.mean((prediction - X) ** 2, axis=1) + assert np.allclose(errors, manual) + assert summary["metric"] == "rmse" + assert summary["n_samples"] == len(X) + + +def test_convlstm_metrics_use_last_frame_as_default_target(): + """ConvLSTM metrics should compare predictions with the final input frame.""" + X = np.random.randn(16, 3, 1, 8, 8).astype("float32") + model = ConvLSTMAutoencoder(k=4, device="cpu") + model.fit(X, **_fit_kwargs()) + + prediction = model.predict(X, batch_size=4, verbose=0) + errors = model.reconstruction_error( + X, + metric="mse", + reduction="sample", + batch_size=4, + ) + manual = np.mean( + (prediction - X[:, -1]) ** 2, + axis=(1, 2, 3), + ) + + assert np.allclose(errors, manual) + + +def test_standard_autoencoder_checkpoint_round_trip(tmp_path): + """A standard autoencoder checkpoint should restore identical predictions.""" + X = np.random.randn(16, 10).astype("float32") + model = StandardAutoencoder(k=3, hidden_dims=[8], device="cpu") + model.fit(X, **_fit_kwargs()) + + checkpoint = tmp_path / "standard_autoencoder.pt" + model.save_pytorch_model(checkpoint) + loaded = StandardAutoencoder.from_pytorch_model( + checkpoint, + device="cpu", + ) + loaded_into_instance = StandardAutoencoder(device="cpu") + loaded_into_instance.load_pytorch_model(checkpoint) + + original = model.predict(X, batch_size=4, verbose=0) + restored = loaded.predict(X, batch_size=4, verbose=0) + restored_instance = loaded_into_instance.predict( + X, + batch_size=4, + verbose=0, + ) + + assert loaded.is_fitted + assert loaded_into_instance.is_fitted + assert np.allclose(original, restored) + assert np.allclose(original, restored_instance) + + +def test_convlstm_checkpoint_round_trip(tmp_path): + """A ConvLSTM checkpoint should restore identical predictions.""" + X = np.random.randn(16, 3, 1, 8, 8).astype("float32") + model = ConvLSTMAutoencoder(k=4, device="cpu") + model.fit(X, **_fit_kwargs()) + + checkpoint = tmp_path / "convlstm_autoencoder.pt" + model.save_pytorch_model(checkpoint) + loaded = ConvLSTMAutoencoder.from_pytorch_model( + checkpoint, + device="cpu", + ) + + original = model.predict(X, batch_size=4, verbose=0) + restored = loaded.predict(X, batch_size=4, verbose=0) + + assert loaded.is_fitted + assert np.allclose(original, restored) diff --git a/tests/deeplearning/test_autoencoders.py b/tests/deeplearning/test_autoencoders.py index d4e7fa3..211315f 100644 --- a/tests/deeplearning/test_autoencoders.py +++ b/tests/deeplearning/test_autoencoders.py @@ -6,10 +6,8 @@ Run from the repository root with: - pytest -q tests/test_deeplearning_autoencoders.py + pytest -q tests/deeplearning/test_autoencoders.py -Some tests are marked xfail because they document likely current bugs in the -implementation. After fixing each issue, remove the corresponding xfail marker. """ import numpy as np