fix: general fixes and improvements

Author: marcpinet

neuralnetlib/callbacks.py

@@ -6,124 +6,131 @@
 
 class ModelWeightManager:
     @staticmethod
-    def get_model_weights(model) -> list[np.ndarray]:
-        """Extract weights from any model type."""
-        weights = []
+    def get_model_weights(model) -> list[tuple[np.ndarray, np.ndarray | None]]:
+        """Extract weights and biases from any model type."""
+        params = []
+
+        def get_params_from_layer(layer):
+            if hasattr(layer, 'weights'):
+                weights = layer.weights.copy()
+                bias = layer.bias.copy() if hasattr(layer, 'bias') else None
+                return (weights, bias)
+            return None
+
+        def get_params_from_dense_layers(layers):
+            layer_params = []
+            for layer in layers:
+                p = get_params_from_layer(layer)
+                if p:
+                    layer_params.append(p)
+            return layer_params
 
         if hasattr(model, 'layers'):  # Sequential model
-            weights.extend(
-                [layer.weights for layer in model.layers if hasattr(layer, 'weights')])
+            for layer in model.layers:
+                p = get_params_from_layer(layer)
+                if p:
+                    params.append(p)
 
         elif hasattr(model, 'encoder_layers') and hasattr(model, 'decoder_layers'):  # Autoencoder
-            weights.extend(
-                [layer.weights for layer in model.encoder_layers if hasattr(layer, 'weights')])
-            weights.extend(
-                [layer.weights for layer in model.decoder_layers if hasattr(layer, 'weights')])
+            for layer in model.encoder_layers:
+                p = get_params_from_layer(layer)
+                if p:
+                    params.append(p)
+            for layer in model.decoder_layers:
+                p = get_params_from_layer(layer)
+                if p:
+                    params.append(p)
 
-        elif hasattr(model, 'embedding'):  # Transformer
-            if hasattr(model.embedding, 'weights'):
-                weights.append(model.embedding.weights)
+        elif hasattr(model, 'src_embedding'):  # Transformer
+            params.append(get_params_from_layer(model.src_embedding))
+            params.append(get_params_from_layer(model.tgt_embedding))
 
             for encoder_layer in model.encoder_layers:
-                if hasattr(encoder_layer, 'attention'):
-                    weights.extend([
-                        encoder_layer.attention.query_dense.weights,
-                        encoder_layer.attention.key_dense.weights,
-                        encoder_layer.attention.value_dense.weights,
-                        encoder_layer.attention.output_dense.weights
-                    ])
-                if hasattr(encoder_layer, 'ffn'):
-                    weights.extend([
-                        encoder_layer.ffn.dense1.weights,
-                        encoder_layer.ffn.dense2.weights
-                    ])
+                params.extend(get_params_from_dense_layers([
+                    encoder_layer.attention.query_dense,
+                    encoder_layer.attention.key_dense,
+                    encoder_layer.attention.value_dense,
+                    encoder_layer.attention.output_dense,
+                    encoder_layer.ffn.dense1,
+                    encoder_layer.ffn.dense2
+                ]))
 
             for decoder_layer in model.decoder_layers:
-                if hasattr(decoder_layer, 'self_attention'):
-                    weights.extend([
-                        decoder_layer.self_attention.query_dense.weights,
-                        decoder_layer.self_attention.key_dense.weights,
-                        decoder_layer.self_attention.value_dense.weights,
-                        decoder_layer.self_attention.output_dense.weights
-                    ])
-                if hasattr(decoder_layer, 'cross_attention'):
-                    weights.extend([
-                        decoder_layer.cross_attention.query_dense.weights,
-                        decoder_layer.cross_attention.key_dense.weights,
-                        decoder_layer.cross_attention.value_dense.weights,
-                        decoder_layer.cross_attention.output_dense.weights
-                    ])
-                if hasattr(decoder_layer, 'ffn'):
-                    weights.extend([
-                        decoder_layer.ffn.dense1.weights,
-                        decoder_layer.ffn.dense2.weights
-                    ])
-
-            if hasattr(model.output_layer, 'weights'):
-                weights.append(model.output_layer.weights)
-
-        return weights
+                params.extend(get_params_from_dense_layers([
+                    decoder_layer.self_attention.query_dense,
+                    decoder_layer.self_attention.key_dense,
+                    decoder_layer.self_attention.value_dense,
+                    decoder_layer.self_attention.output_dense,
+                    decoder_layer.cross_attention.query_dense,
+                    decoder_layer.cross_attention.key_dense,
+                    decoder_layer.cross_attention.value_dense,
+                    decoder_layer.cross_attention.output_dense,
+                    decoder_layer.ffn.dense1,
+                    decoder_layer.ffn.dense2
+                ]))
+
+            params.append(get_params_from_layer(model.output_layer))
+
+        return [p for p in params if p is not None]
 
     @staticmethod
-    def set_model_weights(model, weights: list[np.ndarray]) -> None:
-        """Restore weights to any model type."""
-        weight_idx = 0
+    def set_model_weights(model, params: list[tuple[np.ndarray, np.ndarray | None]]) -> None:
+        """Restore weights and biases to any model type."""
+        param_idx = 0
+
+        def set_params_for_layer(layer):
+            nonlocal param_idx
+            if hasattr(layer, 'weights'):
+                if param_idx < len(params):
+                    weights, bias = params[param_idx]
+                    layer.weights = weights.copy()
+                    if hasattr(layer, 'bias') and bias is not None:
+                        layer.bias = bias.copy()
+                    param_idx += 1
+        
+        def set_params_for_dense_layers(layers):
+            for layer in layers:
+                set_params_for_layer(layer)
 
         if hasattr(model, 'layers'):  # Sequential model
             for layer in model.layers:
-                if hasattr(layer, 'weights'):
-                    layer.weights = weights[weight_idx]
-                    weight_idx += 1
+                set_params_for_layer(layer)
 
         elif hasattr(model, 'encoder_layers') and hasattr(model, 'decoder_layers'):  # Autoencoder
             for layer in model.encoder_layers:
-                if hasattr(layer, 'weights'):
-                    layer.weights = weights[weight_idx]
-                    weight_idx += 1
-
+                set_params_for_layer(layer)
             for layer in model.decoder_layers:
-                if hasattr(layer, 'weights'):
-                    layer.weights = weights[weight_idx]
-                    weight_idx += 1
+                set_params_for_layer(layer)
 
-        elif hasattr(model, 'embedding'):
-            if hasattr(model.embedding, 'weights'):
-                model.embedding.weights = weights[weight_idx]
-                weight_idx += 1
+        elif hasattr(model, 'src_embedding'): # Transformer
+            set_params_for_layer(model.src_embedding)
+            set_params_for_layer(model.tgt_embedding)
 
             for encoder_layer in model.encoder_layers:
-                if hasattr(encoder_layer, 'attention'):
-                    encoder_layer.attention.query_dense.weights = weights[weight_idx]
-                    encoder_layer.attention.key_dense.weights = weights[weight_idx + 1]
-                    encoder_layer.attention.value_dense.weights = weights[weight_idx + 2]
-                    encoder_layer.attention.output_dense.weights = weights[weight_idx + 3]
-                    weight_idx += 4
-                if hasattr(encoder_layer, 'ffn'):
-                    encoder_layer.ffn.dense1.weights = weights[weight_idx]
-                    encoder_layer.ffn.dense2.weights = weights[weight_idx + 1]
-                    weight_idx += 2
+                set_params_for_dense_layers([
+                    encoder_layer.attention.query_dense,
+                    encoder_layer.attention.key_dense,
+                    encoder_layer.attention.value_dense,
+                    encoder_layer.attention.output_dense,
+                    encoder_layer.ffn.dense1,
+                    encoder_layer.ffn.dense2
+                ])
 
             for decoder_layer in model.decoder_layers:
-                if hasattr(decoder_layer, 'self_attention'):
-                    decoder_layer.self_attention.query_dense.weights = weights[weight_idx]
-                    decoder_layer.self_attention.key_dense.weights = weights[weight_idx + 1]
-                    decoder_layer.self_attention.value_dense.weights = weights[weight_idx + 2]
-                    decoder_layer.self_attention.output_dense.weights = weights[weight_idx + 3]
-                    weight_idx += 4
-                if hasattr(decoder_layer, 'cross_attention'):
-                    decoder_layer.cross_attention.query_dense.weights = weights[weight_idx]
-                    decoder_layer.cross_attention.key_dense.weights = weights[weight_idx + 1]
-                    decoder_layer.cross_attention.value_dense.weights = weights[weight_idx + 2]
-                    decoder_layer.cross_attention.output_dense.weights = weights[weight_idx + 3]
-                    weight_idx += 4
-                if hasattr(decoder_layer, 'ffn'):
-                    decoder_layer.ffn.dense1.weights = weights[weight_idx]
-                    decoder_layer.ffn.dense2.weights = weights[weight_idx + 1]
-                    weight_idx += 2
-
-            # Restore output layer weights
-            if hasattr(model.output_layer, 'weights'):
-                model.output_layer.weights = weights[weight_idx]
+                set_params_for_dense_layers([
+                    decoder_layer.self_attention.query_dense,
+                    decoder_layer.self_attention.key_dense,
+                    decoder_layer.self_attention.value_dense,
+                    decoder_layer.self_attention.output_dense,
+                    decoder_layer.cross_attention.query_dense,
+                    decoder_layer.cross_attention.key_dense,
+                    decoder_layer.cross_attention.value_dense,
+                    decoder_layer.cross_attention.output_dense,
+                    decoder_layer.ffn.dense1,
+                    decoder_layer.ffn.dense2
+                ])
+            
+            set_params_for_layer(model.output_layer)
 
 
 class Callback:

neuralnetlib/layers.py

@@ -153,9 +153,8 @@ def forward_pass(self, input_data: np.ndarray) -> np.ndarray:
         self.input_shape = input_data.shape
         self.input = input_data
 
-        if len(input_data.shape) == 1 and self.input_dim:
-            batch_size = input_data.shape[0]
-            input_data = input_data.reshape(batch_size, self.input_dim)
+        if input_data.ndim == 1:
+            input_data = input_data.reshape(1, -1)
             self.input = input_data
 
         if len(input_data.shape) == 3:
@@ -316,13 +315,7 @@ def forward_pass(self, input_data: np.ndarray, training: bool = True) -> np.ndar
     def backward_pass(self, output_error: np.ndarray) -> np.ndarray:
         if self.adaptive:
             return self.dropout_impl.gradient(output_error)
-
-        if output_error.shape[0] != self.mask.shape[0]:
-            rng = np.random.default_rng(
-                self.random_state if self.random_state is not None else int(time.time_ns()))
-            self.mask = rng.binomial(1, 1 - self.rate,
-                                    size=(output_error.shape[0], self.mask.shape[1])) / (1 - self.rate)
-                                    
+        
         return output_error * self.mask
 
     def get_config(self) -> dict:
@@ -373,22 +366,21 @@ def __init__(self, filters: int, kernel_size: int | tuple, strides: int | tuple
 
     def initialize_weights(self, input_shape: tuple):
         _, _, _, in_channels = input_shape
+        fan_in = np.prod(self.kernel_size) * in_channels
+        fan_out = np.prod(self.kernel_size) * self.filters
 
         self.rng = np.random.default_rng(
             self.random_state if self.random_state is not None else int(time.time_ns()))
 
-        if self.weights_init == "xavier":
-            self.weights = self.rng.normal(0, np.sqrt(2 / (np.prod(self.kernel_size) * in_channels)),
-                                           (*self.kernel_size, in_channels, self.filters))
+        if self.weights_init == "glorot_uniform" or self.weights_init == "xavier":
+            limit = np.sqrt(6 / (fan_in + fan_out))
+            self.weights = self.rng.uniform(-limit, limit, (*self.kernel_size, in_channels, self.filters))
         elif self.weights_init == "he":
-            self.weights = self.rng.normal(0, np.sqrt(2 / (in_channels * np.prod(self.kernel_size))),
-                                           (*self.kernel_size, in_channels, self.filters))
+            self.weights = self.rng.normal(0, np.sqrt(2 / fan_in), (*self.kernel_size, in_channels, self.filters))
         elif self.weights_init == "default":
-            self.weights = self.rng.normal(
-                0, 0.01, (*self.kernel_size, in_channels, self.filters))
+            self.weights = self.rng.normal(0, 0.01, (*self.kernel_size, in_channels, self.filters))
         else:
-            raise ValueError(
-                "Invalid weights_init value. Possible values are 'xavier', 'he', and 'default'.")
+            raise ValueError("Invalid weights_init value. Possible values are 'xavier', 'he', and 'default'.")
 
         if self.bias_init == "default":
             self.bias = np.zeros((1, self.filters))

neuralnetlib/losses.py

@@ -81,7 +81,7 @@ def __call__(self, y_true: np.ndarray, y_pred: np.ndarray) -> float:
         return np.mean(np.square(y_true - y_pred))
 
     def derivative(self, y_true: np.ndarray, y_pred: np.ndarray) -> np.ndarray:
-        return 2 * (y_pred - y_true) / y_true.shape[0]
+        return 2 * (y_pred - y_true) / y_true.size
 
     def __str__(self):
         return "MeanSquaredError"

neuralnetlib/metrics.py

@@ -307,18 +307,24 @@ def mean_absolute_percentage_error(y_pred: np.ndarray, y_true: np.ndarray, thres
     return np.mean(np.abs((y_true[mask] - y_pred[mask]) / y_true[mask])) * 100
 
 
-def r2_score(y_pred: np.ndarray, y_true: np.ndarray, threshold: float = 0.5) -> float:
+def r2_score(y_pred: np.ndarray, y_true: np.ndarray) -> float:
     y_pred, y_true = _reshape_inputs(y_pred, y_true)
-    if y_pred.shape[1] == 1:
-        y_pred_classes = (y_pred >= threshold).astype(int).ravel()
-        y_true_classes = y_true.ravel()
-    else:
-        y_pred_classes = np.argmax(y_pred, axis=1)
-        y_true_classes = np.argmax(y_true, axis=1)
 
-    ss_res = np.sum((y_true_classes - y_pred_classes) ** 2)
-    ss_tot = np.sum((y_true_classes - np.mean(y_true_classes)) ** 2)
-    return 1 - (ss_res / ss_tot) if ss_tot != 0 else 0.0
+    if y_pred.shape[1] == 1:
+        y_pred_ = y_pred.ravel()
+        y_true_ = y_true.ravel()
+        ss_res = np.sum((y_true_ - y_pred_) ** 2)
+        ss_tot = np.sum((y_true_ - np.mean(y_true_)) ** 2)
+        return 1.0 - (ss_res / ss_tot) if ss_tot != 0 else 0.0
+
+    r2s = []
+    for j in range(y_pred.shape[1]):
+        yp = y_pred[:, j]
+        yt = y_true[:, j]
+        ss_res = np.sum((yt - yp) ** 2)
+        ss_tot = np.sum((yt - np.mean(yt)) ** 2)
+        r2s.append(1.0 - (ss_res / ss_tot) if ss_tot != 0 else 0.0)
+    return float(np.mean(r2s)) if r2s else 0.0
 
 
 def bleu_score(y_pred: np.ndarray, y_true: np.ndarray, threshold: float | None = None, n_gram: int = 4, smooth: bool = False) -> float:

neuralnetlib/models.py

@@ -193,8 +193,6 @@ def forward_pass(self, X: np.ndarray, training: bool = True, labels: np.ndarray
         return X
 
     def backward_pass(self, error: np.ndarray, gan: bool = False, compute_only: bool = False) -> np.ndarray:
-        if self.n_classes is not None and error.shape[1] > error.shape[1] - self.n_classes:
-            error = error[:, :-self.n_classes]
 
         for i, layer in enumerate(reversed(self.layers)):
             if i == 0 and isinstance(layer, Activation):