!

Author: Ciznia

src/my_torch/network/core.py

@@ -2,6 +2,8 @@
 
 import json
 import random
+from functools import partial
+from multiprocessing import Pool, cpu_count
 from typing import Callable, List, Tuple
 
 import numpy as np
@@ -10,6 +12,33 @@
 from ..optimizer import Optimizer, SGDOptimizer
 
 
+def _compute_neuron_forward(args):
+    """Compute forward pass for a single neuron."""
+    weights, bias, inputs, activation_fn_name, is_softmax = args
+    z = sum(w * inp for w, inp in zip(weights, inputs)) + bias
+    return z
+
+
+def _compute_neuron_gradient(args):
+    """Compute gradients for a single neuron."""
+    neuron_idx, weights, delta, prev_activations = args
+    grad_w = [delta * a for a in prev_activations]
+    grad_b = delta
+    return neuron_idx, grad_w, grad_b
+
+
+def _compute_neuron_backprop(args):
+    """Compute backpropagation delta for a single neuron in hidden layer."""
+    neuron_idx, next_weights, next_deltas, activation, act_deriv = args
+    # Sum contributions from all neurons in next layer
+    s = sum(
+        next_weights[j][neuron_idx] * next_deltas[j]
+        for j in range(len(next_deltas))
+    )
+    delta = s * act_deriv(activation)
+    return neuron_idx, delta
+
+
 class Network:
     def __init__(
         self,
@@ -20,6 +49,7 @@ def __init__(
         optimizer: Optimizer | None = None,
         class_weights: List[float] | None = None,
         seed: int | None = None,
+        n_processes: int | None = None,
     ):
         if len(layer_sizes) < 2:
             raise ValueError("Need at least input and output layer")
@@ -32,6 +62,7 @@ def __init__(
         self.loss_fn = loss
         self.class_weights = class_weights
         self.optimizer = optimizer or SGDOptimizer(learning_rate=0.3)
+        self.n_processes = n_processes or max(1, cpu_count() - 1)
 
         if output_activation == "softmax" and layer_sizes[-1] == 1:
             raise ValueError("Softmax requires >1 output neurons")
@@ -93,7 +124,6 @@ def evaluate(
 
         return correct / len(dataset)
 
-    # Method stubs for dynamically attached methods (implemented in submodules)
     @staticmethod
     def get_activation(name: str) -> Callable[[float], float]:
         """Get activation function by name."""
@@ -115,75 +145,73 @@ def get_activation_derivative(name: str) -> Callable[[float], float]:
         }[name]
 
     def forward(
-        self, inputs: List[float]
+        self, inputs: List[float], pool: Pool | None = None
     ) -> Tuple[List[List[float]], List[List[float]]]:
-        """Forward pass through the network."""
+        """Forward pass through the network with optional multiprocessing."""
         if len(inputs) != self.layer_sizes[0]:
             raise ValueError(
                 f"Input size {len(inputs)} doesn't match network input {self.layer_sizes[0]}"
             )
 
-        # Convert inputs to numpy array
-        a = np.array(inputs)
-
-        # Store activations and z values for each layer
-        activations_list = [a]
+        activations_list = [inputs]
         zs = []
+        current_input = inputs
 
         for layer_idx in range(len(self.weights)):
-            # Extract weight matrix and bias vector for current layer
             w_mat = self.weights[layer_idx]
             b_vec = self.biases[layer_idx]
+            is_output = layer_idx == len(self.weights) - 1
+            is_softmax = is_output and self.output_activation == "softmax"
+
+            # Parallel computation of z values for all neurons in this layer
+            if (
+                pool is not None and len(w_mat) > 10
+            ):  # Use MP for larger layers
+                args = [
+                    (
+                        w_mat[i],
+                        b_vec[i],
+                        current_input,
+                        (
+                            self.output_activation
+                            if is_output
+                            else self.hidden_activation
+                        ),
+                        is_softmax,
+                    )
+                    for i in range(len(w_mat))
+                ]
+                z_layer = pool.map(_compute_neuron_forward, args)
+            else:
+                z_layer = [
+                    sum(w * inp for w, inp in zip(w_mat[i], current_input))
+                    + b_vec[i]
+                    for i in range(len(w_mat))
+                ]
 
-            # Compute z values: z = W * a + b
-            z_layer = np.dot(w_mat, a) + b_vec
             zs.append(z_layer)
 
-            # Determine if it's the output layer
-            is_output = layer_idx == len(self.weights) - 1
-
-            # Activation function for output layer (e.g., softmax)
-            if is_output and self.output_activation == "softmax":
-                a_next = self.softmax(z_layer)
+            # Apply activation function
+            if is_softmax:
+                z_array = np.array(z_layer)
+                exp_z = np.exp(z_array - np.max(z_array))
+                a_next = (exp_z / np.sum(exp_z)).tolist()
             else:
-                # Use appropriate activation function for hidden layers
                 act_fn = self.get_activation(
                     self.output_activation
                     if is_output
                     else self.hidden_activation
                 )
-                a_next = act_fn(z_layer)
+                a_next = [act_fn(z) for z in z_layer]
 
             activations_list.append(a_next)
-            a = a_next
-
-        # Convert activations back to list of lists if needed
-        activations_list = [a.tolist() for a in activations_list]
-        zs = [z.tolist() for z in zs]
+            current_input = a_next
 
         return activations_list, zs
 
-    def softmax(self, z: np.ndarray) -> np.ndarray:
-        """Softmax activation function."""
-        exp_z = np.exp(z - np.max(z))  # Numerical stability
-        return exp_z / np.sum(exp_z, axis=-1, keepdims=True)
-
-    def get_activation(self, activation_type: str):
-        """Returns the activation function based on the type."""
-        if activation_type == "sigmoid":
-            return lambda x: 1 / (1 + np.exp(-x))
-        elif activation_type == "tanh":
-            return np.tanh
-        elif activation_type == "relu":
-            return lambda x: np.maximum(0, x)
-        else:
-            raise ValueError(
-                f"Activation function {activation_type} not recognized"
-            )
-
     def predict(self, inputs: List[float]) -> List[float]:
         """Make a prediction for given inputs."""
-        activations_list, _ = self.forward(inputs)
+        activations_list, _ = self.forward(inputs, pool=None)
         return activations_list[-1]
 
     def compute_output_delta(
@@ -212,52 +240,102 @@ def compute_output_delta(
             ]
 
     def compute_hidden_deltas(
-        self, activations_list: List[List[float]], output_delta: List[float]
+        self,
+        activations_list: List[List[float]],
+        output_delta: List[float],
+        pool: Pool | None = None,
     ) -> List[List[float]]:
-        """Backpropagate delta through hidden layers."""
+        """Backpropagate delta through hidden layers with multiprocessing."""
         deltas: List[List[float]] = [
             [] for _ in range(len(self.layer_sizes) - 1)
         ]
         last_layer_idx = len(self.layer_sizes) - 2
         deltas[last_layer_idx] = output_delta
 
         act_deriv = self.get_activation_derivative(self.hidden_activation)
+
         for layer_idx in range(last_layer_idx - 1, -1, -1):
-            layer_deltas = []
-            for i in range(self.layer_sizes[layer_idx + 1]):
-                s = sum(
-                    self.weights[layer_idx + 1][j][i]
-                    * deltas[layer_idx + 1][j]
-                    for j in range(self.layer_sizes[layer_idx + 2])
-                )
-                a_val = activations_list[layer_idx + 1][i]
-                layer_deltas.append(s * act_deriv(a_val))
+            layer_size = self.layer_sizes[layer_idx + 1]
+            next_weights = self.weights[layer_idx + 1]
+            next_deltas = deltas[layer_idx + 1]
+            layer_activations = activations_list[layer_idx + 1]
+
+            # Parallel computation of deltas for all neurons in this layer
+            if pool is not None and layer_size > 10:
+                args = [
+                    (
+                        i,
+                        next_weights,
+                        next_deltas,
+                        layer_activations[i],
+                        act_deriv,
+                    )
+                    for i in range(layer_size)
+                ]
+                results = pool.map(_compute_neuron_backprop, args)
+                # Sort by neuron index to maintain order
+                results.sort(key=lambda x: x[0])
+                layer_deltas = [delta for _, delta in results]
+            else:
+                layer_deltas = []
+                for i in range(layer_size):
+                    s = sum(
+                        next_weights[j][i] * next_deltas[j]
+                        for j in range(len(next_deltas))
+                    )
+                    layer_deltas.append(s * act_deriv(layer_activations[i]))
+
             deltas[layer_idx] = layer_deltas
 
         return deltas
 
     def compute_gradients(
-        self, activations_list: List[List[float]], deltas: List[List[float]]
+        self,
+        activations_list: List[List[float]],
+        deltas: List[List[float]],
+        pool: Pool | None = None,
     ) -> Tuple[List[List[List[float]]], List[List[float]]]:
-        """Compute weight and bias gradients from deltas."""
+        """Compute weight and bias gradients with multiprocessing."""
         grad_w = [
             [[0.0 for _ in row] for row in layer] for layer in self.weights
         ]
         grad_b = [[0.0 for _ in layer] for layer in self.biases]
 
         for layer_idx in range(len(self.weights)):
-            for neuron_idx in range(len(self.weights[layer_idx])):
-                for w_idx in range(len(self.weights[layer_idx][neuron_idx])):
-                    grad_w[layer_idx][neuron_idx][w_idx] = (
-                        deltas[layer_idx][neuron_idx]
-                        * activations_list[layer_idx][w_idx]
+            layer_size = len(self.weights[layer_idx])
+            prev_activations = activations_list[layer_idx]
+
+            # Parallel computation of gradients for all neurons in this layer
+            if pool is not None and layer_size > 10:
+                args = [
+                    (
+                        i,
+                        self.weights[layer_idx][i],
+                        deltas[layer_idx][i],
+                        prev_activations,
                     )
-                grad_b[layer_idx][neuron_idx] = deltas[layer_idx][neuron_idx]
+                    for i in range(layer_size)
+                ]
+                results = pool.map(_compute_neuron_gradient, args)
+
+                for neuron_idx, gw, gb in results:
+                    grad_w[layer_idx][neuron_idx] = gw
+                    grad_b[layer_idx][neuron_idx] = gb
+            else:
+                for neuron_idx in range(layer_size):
+                    delta = deltas[layer_idx][neuron_idx]
+                    grad_w[layer_idx][neuron_idx] = [
+                        delta * a for a in prev_activations
+                    ]
+                    grad_b[layer_idx][neuron_idx] = delta
 
         return grad_w, grad_b
 
     def backprop_sample(
-        self, activations_list: List[List[float]], target: List[float]
+        self,
+        activations_list: List[List[float]],
+        target: List[float],
+        pool: Pool | None = None,
     ) -> Tuple[List[List[List[float]]], List[List[float]], float, bool]:
         """Full backpropagation for a single sample."""
         output = activations_list[-1]
@@ -275,10 +353,12 @@ def backprop_sample(
 
         # Compute deltas
         output_delta = self.compute_output_delta(activations_list, target)
-        deltas = self.compute_hidden_deltas(activations_list, output_delta)
+        deltas = self.compute_hidden_deltas(
+            activations_list, output_delta, pool
+        )
 
         # Compute gradients
-        grad_w, grad_b = self.compute_gradients(activations_list, deltas)
+        grad_w, grad_b = self.compute_gradients(activations_list, deltas, pool)
 
         return grad_w, grad_b, loss, is_correct
 
@@ -294,40 +374,43 @@ def train_epoch(
         total_correct = 0
         random.shuffle(dataset)
 
-        for start in range(0, len(dataset), batch_size):
-            batch = dataset[start : start + batch_size]
+        # Create process pool for the entire epoch
+        with Pool(processes=self.n_processes) as pool:
+            for start in range(0, len(dataset), batch_size):
+                batch = dataset[start : start + batch_size]
 
-            acc_grad_w = [
-                [[0.0 for _ in row] for row in layer] for layer in self.weights
-            ]
-            acc_grad_b = [[0.0 for _ in layer] for layer in self.biases]
+                acc_grad_w = [
+                    [[0.0 for _ in row] for row in layer]
+                    for layer in self.weights
+                ]
+                acc_grad_b = [[0.0 for _ in layer] for layer in self.biases]
 
-            for x, target in batch:
-                activations_list, _ = self.forward(x)
-                grad_w, grad_b, loss, is_correct = self.backprop_sample(
-                    activations_list, target
-                )
+                for x, target in batch:
+                    activations_list, _ = self.forward(x, pool)
+                    grad_w, grad_b, loss, is_correct = self.backprop_sample(
+                        activations_list, target, pool
+                    )
+
+                    total_loss += loss
+                    if is_correct:
+                        total_correct += 1
 
-                total_loss += loss
-                if is_correct:
-                    total_correct += 1
+                    for li in range(len(self.weights)):
+                        for ni in range(len(self.weights[li])):
+                            for wi in range(len(self.weights[li][ni])):
+                                acc_grad_w[li][ni][wi] += grad_w[li][ni][wi]
+                            acc_grad_b[li][ni] += grad_b[li][ni]
 
+                bsz = len(batch)
                 for li in range(len(self.weights)):
                     for ni in range(len(self.weights[li])):
                         for wi in range(len(self.weights[li][ni])):
-                            acc_grad_w[li][ni][wi] += grad_w[li][ni][wi]
-                        acc_grad_b[li][ni] += grad_b[li][ni]
-
-            bsz = len(batch)
-            for li in range(len(self.weights)):
-                for ni in range(len(self.weights[li])):
-                    for wi in range(len(self.weights[li][ni])):
-                        acc_grad_w[li][ni][wi] /= bsz
-                    acc_grad_b[li][ni] /= bsz
-
-            self.optimizer.update(
-                self.weights, self.biases, acc_grad_w, acc_grad_b
-            )
+                            acc_grad_w[li][ni][wi] /= bsz
+                        acc_grad_b[li][ni] /= bsz
+
+                self.optimizer.update(
+                    self.weights, self.biases, acc_grad_w, acc_grad_b
+                )
 
         self.optimizer.decay_lr()