# @Authored by Claude Sonnet 4.6, Co-Authored by Sujith M A, Created 2026-05-31, Last Updated 2026-05-31 import logging import math from typing import List import numpy as np import torch import torch.nn as nn from sklearn.preprocessing import LabelEncoder from torch.utils.data import DataLoader, TensorDataset from core.data_bundle import DataBundle from models.base import BaseModel logger = logging.getLogger(__name__) ACTIVATION_MAP: dict = { "relu": nn.ReLU, "tanh": nn.Tanh, "sigmoid": nn.Sigmoid, "leaky_relu": nn.LeakyReLU, "elu": nn.ELU, } def _build_fc_block( input_size: int, output_size: int, activation_name: str, dropout_rate: float, ) -> List[nn.Module]: """Returns [Linear, Activation, Dropout] layers for one FC block.""" activation_class = ACTIVATION_MAP.get(activation_name, nn.ReLU) return [ nn.Linear(input_size, output_size), activation_class(), nn.Dropout(p=dropout_rate), ] def _build_fcnn( input_size: int, hidden_layers: List[int], num_outputs: int, activation_name: str, dropout_rate: float, ) -> nn.Sequential: """Build a fully-connected feedforward network from config.""" layer_sizes = [input_size] + hidden_layers layers: List[nn.Module] = [] for layer_index in range(len(layer_sizes) - 1): layers.extend( _build_fc_block( layer_sizes[layer_index], layer_sizes[layer_index + 1], activation_name, dropout_rate, ) ) layers.append(nn.Linear(layer_sizes[-1], num_outputs)) return nn.Sequential(*layers) class PyTorchFCNNClassifierWrapper(BaseModel): """Fully-connected neural network classifier built from config.yaml architecture.""" def __init__(self, config: dict) -> None: super().__init__(config) self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.label_encoder: LabelEncoder = LabelEncoder() def train(self, data: DataBundle) -> None: merged_config = self._merge_hyperparameter_overrides(data) X_train = data.common.X_train.astype(np.float32) # Remap arbitrary class labels to consecutive integers 0..N-1. # CrossEntropyLoss requires target indices in [0, num_classes); without this # remapping a dataset with labels {0,1,3,...,14} causes "Target out of bounds". y_train = self.label_encoder.fit_transform(data.common.y_train).astype(np.int64) num_features = X_train.shape[1] # Derive num_classes from actual data so the model output layer matches the # dataset even when the config default (e.g. 10) doesn't match class count. num_classes = len(self.label_encoder.classes_) hidden_layers: List[int] = list(merged_config["hidden_layers"]) activation_name: str = merged_config["activation"] dropout_rate: float = float(merged_config["dropout"]) learning_rate: float = float(merged_config["learning_rate"]) num_epochs: int = int(merged_config["epochs"]) batch_size: int = int(merged_config["batch_size"]) weight_decay: float = float(merged_config.get("weight_decay", 1e-4)) self.model = _build_fcnn( num_features, hidden_layers, num_classes, activation_name, dropout_rate ).to(self.device) optimizer = torch.optim.Adam( self.model.parameters(), lr=learning_rate, weight_decay=weight_decay ) loss_fn = nn.CrossEntropyLoss() X_tensor = torch.tensor(X_train) y_tensor = torch.tensor(y_train) dataset = TensorDataset(X_tensor, y_tensor) loader = DataLoader(dataset, batch_size=batch_size, shuffle=True) self.model.train() for epoch_index in range(num_epochs): epoch_loss = 0.0 for X_batch, y_batch in loader: X_batch = X_batch.to(self.device) y_batch = y_batch.to(self.device) optimizer.zero_grad() logits = self.model(X_batch) loss = loss_fn(logits, y_batch) loss.backward() optimizer.step() epoch_loss += loss.item() logger.debug( "=> PyTorchFCNNClassifier epoch %d/%d loss=%.4f", epoch_index + 1, num_epochs, epoch_loss / len(loader), ) def predict(self, X: np.ndarray) -> np.ndarray: X_tensor = torch.tensor(X.astype(np.float32)).to(self.device) self.model.eval() with torch.no_grad(): logits = self.model(X_tensor) encoded_predictions = logits.argmax(dim=1).cpu().numpy() # Decode back to original class labels. return self.label_encoder.inverse_transform(encoded_predictions) def _prepare_for_serialization(self) -> nn.Module: return self.model.cpu() class _CNNClassifierNet(nn.Module): """Internal CNN network for classification. Supports Conv1d and Conv2d.""" def __init__( self, conv_dim: int, in_channels: int, conv_channels: List[int], kernel_size: int, fc_input_size: int, fc_layers: List[int], num_classes: int, dropout_rate: float, ) -> None: super().__init__() conv_layer_class = nn.Conv2d if conv_dim == 2 else nn.Conv1d pool_layer_class = nn.MaxPool2d if conv_dim == 2 else nn.MaxPool1d conv_blocks: List[nn.Module] = [] current_channels = in_channels for out_channels in conv_channels: conv_blocks += [ conv_layer_class(current_channels, out_channels, kernel_size, padding=1), nn.ReLU(), pool_layer_class(2), ] current_channels = out_channels self.conv_layers = nn.Sequential(*conv_blocks) fc_sizes = [fc_input_size] + fc_layers fc_blocks: List[nn.Module] = [nn.Flatten()] for layer_index in range(len(fc_sizes) - 1): fc_blocks += [ nn.Linear(fc_sizes[layer_index], fc_sizes[layer_index + 1]), nn.ReLU(), nn.Dropout(p=dropout_rate), ] fc_blocks.append(nn.Linear(fc_sizes[-1], num_classes)) self.fc_layers = nn.Sequential(*fc_blocks) def forward(self, input_tensor: torch.Tensor) -> torch.Tensor: features = self.conv_layers(input_tensor) return self.fc_layers(features) def _compute_cnn_fc_input_size( conv_dim: int, input_spatial_size: int, conv_channels: List[int], kernel_size: int, num_conv_layers: int, ) -> int: """Compute the flattened size after all conv+pool layers.""" spatial_size = input_spatial_size for _ in range(num_conv_layers): # conv with padding=1 preserves size; pool with kernel=2 halves it spatial_size = spatial_size // 2 last_channels = conv_channels[-1] if conv_dim == 2: return last_channels * spatial_size * spatial_size return last_channels * spatial_size class PyTorchCNNClassifierWrapper(BaseModel): """CNN classifier. conv_dim=2 for image input (MNIST), conv_dim=1 for 1D signals.""" def __init__(self, config: dict) -> None: super().__init__(config) self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") def _reshape_for_cnn( self, X: np.ndarray, conv_dim: int ) -> np.ndarray: """Reshape flat (N, features) input into CNN-compatible shape.""" num_samples = X.shape[0] num_features = X.shape[1] if conv_dim == 2: # Infer spatial dimension from sqrt of feature count spatial_side = int(math.isqrt(num_features)) if spatial_side * spatial_side != num_features: # Pad to next perfect square next_square = (spatial_side + 1) ** 2 X = np.pad(X, ((0, 0), (0, next_square - num_features))) spatial_side = spatial_side + 1 logger.debug( "=> PyTorchCNNClassifier: padded features to %d for 2D reshape.", next_square, ) return X.reshape(num_samples, 1, spatial_side, spatial_side) else: # Conv1d: treat each feature as one time step return X.reshape(num_samples, 1, num_features) def train(self, data: DataBundle) -> None: merged_config = self._merge_hyperparameter_overrides(data) X_train = data.common.X_train.astype(np.float32) y_train = data.common.y_train.astype(np.int64) conv_dim: int = int(merged_config["conv_dim"]) conv_channels: List[int] = list(merged_config["conv_channels"]) kernel_size: int = int(merged_config["kernel_size"]) fc_layers: List[int] = list(merged_config["fc_layers"]) num_classes: int = int(merged_config["num_classes"]) dropout_rate: float = float(merged_config["dropout"]) learning_rate: float = float(merged_config["learning_rate"]) num_epochs: int = int(merged_config["epochs"]) batch_size: int = int(merged_config["batch_size"]) weight_decay: float = float(merged_config.get("weight_decay", 1e-4)) X_reshaped = self._reshape_for_cnn(X_train, conv_dim) spatial_size = X_reshaped.shape[2] # height (or length for 1D) fc_input_size = _compute_cnn_fc_input_size( conv_dim, spatial_size, conv_channels, kernel_size, len(conv_channels) ) self.model = _CNNClassifierNet( conv_dim=conv_dim, in_channels=1, conv_channels=conv_channels, kernel_size=kernel_size, fc_input_size=fc_input_size, fc_layers=fc_layers, num_classes=num_classes, dropout_rate=dropout_rate, ).to(self.device) optimizer = torch.optim.Adam( self.model.parameters(), lr=learning_rate, weight_decay=weight_decay ) loss_fn = nn.CrossEntropyLoss() X_tensor = torch.tensor(X_reshaped) y_tensor = torch.tensor(y_train) dataset = TensorDataset(X_tensor, y_tensor) loader = DataLoader(dataset, batch_size=batch_size, shuffle=True) self.model.train() for epoch_index in range(num_epochs): epoch_loss = 0.0 for X_batch, y_batch in loader: X_batch = X_batch.to(self.device) y_batch = y_batch.to(self.device) optimizer.zero_grad() logits = self.model(X_batch) loss = loss_fn(logits, y_batch) loss.backward() optimizer.step() epoch_loss += loss.item() logger.debug( "=> PyTorchCNNClassifier epoch %d/%d loss=%.4f", epoch_index + 1, num_epochs, epoch_loss / len(loader), ) # Store conv_dim so predict() can apply the same reshape self.trained_conv_dim = conv_dim def predict(self, X: np.ndarray) -> np.ndarray: X_reshaped = self._reshape_for_cnn(X.astype(np.float32), self.trained_conv_dim) X_tensor = torch.tensor(X_reshaped).to(self.device) self.model.eval() with torch.no_grad(): logits = self.model(X_tensor) predicted_classes = logits.argmax(dim=1).cpu().numpy() return predicted_classes def _prepare_for_serialization(self) -> nn.Module: return self.model.cpu()