[ENH] Start the networks module with EEGNet example

aeon_neuro/networks/__init__.py

@@ -0,0 +1,6 @@
+"""Deep learning networks."""
+
+__all__ = ["BaseDeepLearningNetwork", "EEGNetNetwork"]
+
+from aeon_neuro.networks._eegnet import EEGNetNetwork
+from aeon_neuro.networks.base import BaseDeepLearningNetwork

aeon_neuro/networks/_eegnet.py

@@ -0,0 +1,204 @@
+"""EEG Network (EEGNet)."""
+
+__all__ = ["EEGNetNetwork"]
+__maintainer__ = ["hadifawaz1999"]
+
+from aeon_neuro.networks.base import BaseDeepLearningNetwork
+
+
+class EEGNetNetwork(BaseDeepLearningNetwork):
+    """Establish the network structure of EEGNet.
+
+    EEGNet [1]_ is a Convolutional Neural Network (CNN) that uses
+    DepthWise Seprabrable Convolutions in order to avoid high
+    number of trainable parameters. EEGNet uses two dimensional
+    convolutions in order to take into consideration channel
+    correlation. The first phase of this model is a temporal
+    standard convolution, in 2D to avoid aggregating channel
+    wise outputs, this is followed by spatial 2D convolution
+    to learn correlations between channels. This is followed by
+    multiple 2D Separable Convolutions (default to 1 layer)
+    with auto-calculated parameters. The final layer applies
+    a dropout and flattens all axes to a single vector.
+
+    Parameters
+    ----------
+    n_temporal_conv_filters : int, default = 8
+        The number of standard convolution filters learned.
+    kernel_size : int, default = 64
+        The length of every convolution filter applied
+        on the temporal axis.
+    depth_multiplier : int, default = 2
+        The number of filters to learn per channel in the
+        spatial convolution phase.
+    dropout_rate : float, default = 0.5
+        The dropout rate to turn off a percentage of neurons
+        to avoid overfitting.
+    n_separable_convolution_layers: int, default = 1,
+        The number of Depthwise Separable Convolution layers
+        applied after the first spatio-temporal convolution
+        phase. Number of filters of these layers is set
+        to the n_temporal_conv_filters*depth_multiplier
+        and the kernel size to kernel_size/4.
+    activation : str, default = "elu"
+        The default activation function used after every
+        convolution block.
+    pool_size : int, default = 4
+        The size of the temporal average pooling operation.
+        In the separable convolution layers, the pool size
+        value is doubled.
+    """
+
+    def __init__(
+        self,
+        n_temporal_conv_filters: int = 8,
+        kernel_size: int = 64,
+        depth_multiplier: int = 2,
+        dropout_rate: float = 0.5,
+        n_separable_convolution_layers: int = 1,
+        activation: str = "elu",
+        pool_size: int = 4,
+    ):
+
+        self.n_temporal_conv_filters = n_temporal_conv_filters
+        self.kernel_size = kernel_size
+        self.depth_multiplier = depth_multiplier
+        self.dropout_rate = dropout_rate
+        self.n_separable_convolution_layers = n_separable_convolution_layers
+        self.activation = activation
+        self.pool_size = pool_size
+
+        super().__init__()
+
+    def _spatio_temporal_convolution_layer(
+        self,
+        input_tensor,
+        n_input_channels,
+        n_filters,
+        kernel_size,
+        strides,
+        padding,
+        activation,
+        use_bias,
+        depth_multiplier,
+    ):
+        import tensorflow as tf
+
+        temporal_conv = tf.keras.layers.Conv2D(
+            filters=n_filters,
+            kernel_size=(kernel_size, 1),
+            strides=strides,
+            padding=padding,
+            use_bias=use_bias,
+        )(input_tensor)
+
+        temporal_conv = tf.keras.layers.BatchNormalization()(temporal_conv)
+
+        spatial_conv = tf.keras.layers.DepthwiseConv2D(
+            kernel_size=(1, n_input_channels),
+            use_bias=use_bias,
+            depth_multiplier=depth_multiplier,
+            depthwise_constraint=tf.keras.constraints.max_norm(1.0),
+        )(temporal_conv)
+
+        spatial_conv = tf.keras.layers.BatchNormalization()(spatial_conv)
+        spatial_conv = tf.keras.layers.Activation(activation=activation)(spatial_conv)
+
+        return spatial_conv
+
+    def _separable_convolution(
+        self,
+        input_tensor,
+        n_filters,
+        kernel_size,
+        strides,
+        padding,
+        use_bias,
+        activation,
+        pool_size,
+        dropout_rate,
+    ):
+        import tensorflow as tf
+
+        conv = tf.keras.layers.SeparableConv2D(
+            filters=n_filters,
+            kernel_size=(kernel_size, 1),
+            padding=padding,
+            strides=strides,
+            use_bias=use_bias,
+        )(input_tensor)
+        conv = tf.keras.layers.BatchNormalization()(conv)
+        conv = tf.keras.layers.Activation(activation=activation)(conv)
+
+        pool = tf.keras.layers.AveragePooling2D(pool_size=(pool_size, 1))(conv)
+
+        dropout = tf.keras.layers.Dropout(dropout_rate)(pool)
+
+        return dropout
+
+    def build_network(self, input_shape, **kwargs):
+        """Construct a network and return its input and output layers.
+
+        Parameters
+        ----------
+        input_shape : tuple
+            The shape of the data fed into the input layer. This function
+            assumes the input_shape is (n_timepoints, n_channels) as it
+            is the tensorflow-keras expectation.
+
+        Returns
+        -------
+        input_layer : a keras layer
+        output_layer : a keras layer
+        """
+        import tensorflow as tf
+
+        input_layer = tf.keras.layers.Input(input_shape)
+
+        n_timepoints = input_shape[0]
+        n_channels = input_shape[1]
+
+        reshape_layer = tf.keras.layers.Reshape(
+            target_shape=(n_timepoints, n_channels, 1)
+        )(input_layer)
+
+        spatio_temporal_convolution = self._spatio_temporal_convolution_layer(
+            input_tensor=reshape_layer,
+            n_input_channels=n_channels,
+            n_filters=self.n_temporal_conv_filters,
+            kernel_size=self.kernel_size,
+            strides=1,
+            padding="same",
+            activation=self.activation,
+            use_bias=False,
+            depth_multiplier=self.depth_multiplier,
+        )
+
+        average_pooling = tf.keras.layers.AveragePooling2D(
+            pool_size=(self.pool_size, 1)
+        )(spatio_temporal_convolution)
+
+        dropout = tf.keras.layers.Dropout(self.dropout_rate)(average_pooling)
+
+        new_n_channels = self.n_temporal_conv_filters * self.depth_multiplier
+
+        x = dropout
+
+        for _ in range(self.n_separable_convolution_layers):
+            x = self._separable_convolution(
+                input_tensor=x,
+                n_filters=new_n_channels,
+                kernel_size=self.kernel_size // 4,
+                strides=1,
+                padding="same",
+                use_bias=False,
+                activation=self.activation,
+                pool_size=self.pool_size * 2,
+                dropout_rate=self.dropout_rate,
+            )
+
+        flatten = tf.keras.layers.Flatten()(x)
+
+        output_layer = tf.keras.layers.Dropout(self.dropout_rate)(flatten)
+
+        return input_layer, output_layer

aeon_neuro/networks/base.py

@@ -0,0 +1,47 @@
+"""Abstract base class for deep learning networks."""
+
+__maintainer__ = ["hadifawaz1999"]
+
+from abc import ABC, abstractmethod
+
+from aeon.utils.repr import get_unchanged_and_required_params_as_str
+from aeon.utils.validation._dependencies import (
+    _check_soft_dependencies,
+)
+
+
+class BaseDeepLearningNetwork(ABC):
+    """Abstract base class for deep learning networks."""
+
+    _config = {
+        "python_dependencies": ["tensorflow"],
+        "structure": "encoder",
+    }
+
+    @abstractmethod
+    def __init__(self, soft_dependencies="tensorflow"):
+        _check_soft_dependencies(soft_dependencies)
+        super().__init__()
+
+    def __repr__(self):
+        """Format str output like scikit-learn estimators."""
+        changed_params = get_unchanged_and_required_params_as_str(self)
+        return f"{self.__class__.__name__}({changed_params})"
+
+    @abstractmethod
+    def build_network(self, input_shape, **kwargs):
+        """Construct a network and return its input and output layers.
+
+        Parameters
+        ----------
+        input_shape : tuple
+            The shape of the data fed into the input layer This function
+            assumes the input_shape is (n_timepoints, n_channels) as it
+            is the tensorflow-keras expectation.
+
+        Returns
+        -------
+        input_layer : a keras layer
+        output_layer : a keras layer
+        """
+        ...

aeon_neuro/networks/tests/__init__.py

@@ -0,0 +1 @@
+"""Deep learning networks tests."""

aeon_neuro/networks/tests/test_eegnet.py

@@ -0,0 +1,29 @@
+"""Tests for EEGNet."""
+
+import pytest
+from aeon.utils.validation._dependencies import _check_soft_dependencies
+
+from aeon_neuro.networks import EEGNetNetwork
+
+
+@pytest.mark.skipif(
+    not _check_soft_dependencies(["tensorflow"], severity="none"),
+    reason="Tensorflow soft dependency unavailable.",
+)
+def test_output_shapes_eegnet():
+    """Testing output shapes of EEGNet."""
+    input_shape = (10, 2)
+    network = EEGNetNetwork(kernel_size=4, pool_size=1)
+
+    # in the case of input shape (10, 2) and kernel_size=4
+    # and pool_size=1, the expected "theoretical" output shape
+    # of this netwroks should be a vector of length 80
+
+    input_layer, output_layer = network.build_network(input_shape=input_shape)
+
+    assert input_layer is not None
+    assert output_layer is not None
+    assert input_layer.shape[1] == input_shape[0]
+    assert input_layer.shape[2] == input_shape[1]
+    assert len(output_layer.shape[1:]) == 1
+    assert int(output_layer.shape[1]) == 80

aeon_neuro/networks/tests/test_network_base.py

@@ -0,0 +1,43 @@
+"""Base for deep learning network base."""
+
+import pytest
+from aeon.testing.data_generation import make_example_3d_numpy
+from aeon.utils.validation._dependencies import _check_soft_dependencies
+
+from aeon_neuro.networks import BaseDeepLearningNetwork
+
+__maintainer__ = ["hadifawaz1999"]
+
+
+class DummyDeepNetwork(BaseDeepLearningNetwork):
+    """A Dummy Deep Network for testing empty base network class save utilities."""
+
+    def __init__(self):
+        super().__init__()
+
+    def build_network(self, input_shape, **kwargs):
+        """Build a neural network."""
+        import tensorflow as tf
+
+        input_layer = tf.keras.layers.Input(input_shape)
+        flatten_layer = tf.keras.layers.Flatten()(input_layer)
+        output_layer = tf.keras.layers.Dense(units=10)(flatten_layer)
+
+        return input_layer, output_layer
+
+
+@pytest.mark.skipif(
+    not _check_soft_dependencies("tensorflow", severity="none"),
+    reason="skip test if required soft dependency not available",
+)
+def test_dummy_deep_network():
+    """Test the dummy network."""
+    dummy_network = DummyDeepNetwork()
+
+    X, y = make_example_3d_numpy()
+
+    input_layer, output_layer = dummy_network.build_network(input_shape=X.shape)
+
+    assert input_layer is not None
+    assert output_layer is not None
+    assert output_layer.shape[-1] == 10

docs/api_reference.md

@@ -14,4 +14,5 @@ api_reference/data_scripts
 api_reference/examples
 api_reference/transformations
 api_reference/utils
+api_reference/networks
 ```

docs/api_reference/networks.rst

@@ -0,0 +1,18 @@
+.. _networks_ref:
+
+Deep learning networks
+======================
+
+``aeon-neuro`` networks are the models behind the deep learning estimators, and can be used in
+their own right to construct bespoke solutions.
+
+
+.. currentmodule:: aeon_neuro.networks
+
+
+.. autosummary::
+    :toctree: auto_generated/
+    :template: class.rst
+
+    BaseDeepLearningNetwork
+    EEGNetNetwork

pyproject.toml

@@ -52,6 +52,10 @@ dependencies = [
 all_extras = [
     "aeon[all_extras]",
 ]
+dl = [
+    "tensorflow>=2.14",
+    "keras>=3.6.0",
+]
 dev = [
     "aeon[dev]",
 ]