ezmsg-learn 1.0__py3-none-any.whl

ezmsg/learn/process/rnn.py

@@ -0,0 +1,266 @@
+import typing
+
+import ezmsg.core as ez
+import numpy as np
+import torch
+from ezmsg.sigproc.base import BaseAdaptiveTransformer, BaseAdaptiveTransformerUnit
+from ezmsg.sigproc.sampler import SampleMessage
+from ezmsg.sigproc.util.profile import profile_subpub
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
+from .base import ModelInitMixin
+from .torch import (
+    TorchModelSettings,
+    TorchModelState,
+    TorchProcessorMixin,
+)
+
+
+class RNNSettings(TorchModelSettings):
+    model_class: str = "ezmsg.learn.model.rnn.RNNModel"
+    """
+    Fully qualified class path of the model to be used.
+    This should be "ezmsg.learn.model.rnn.RNNModel" for this.
+    """
+    reset_hidden_on_fit: bool = True
+    """
+    Whether to reset the hidden state on each fit call.
+    If True, the hidden state will be reset to zero after each fit.
+    If False, the hidden state will be maintained across fit calls.
+    """
+    preserve_state_across_windows: bool | typing.Literal["auto"] = "auto"
+    """
+    Whether to preserve the hidden state across windows.
+    If True, the hidden state will be preserved across windows.
+    If False, the hidden state will be reset at the start of each window.
+    If "auto", preserve if there is no overlap in time windows, otherwise reset.
+    """
+
+
+class RNNState(TorchModelState):
+    hx: typing.Optional[torch.Tensor] = None
+
+
+class RNNProcessor(
+    BaseAdaptiveTransformer[RNNSettings, AxisArray, AxisArray, RNNState],
+    TorchProcessorMixin,
+    ModelInitMixin,
+):
+    def _infer_output_sizes(
+        self, model: torch.nn.Module, n_input: int
+    ) -> dict[str, int]:
+        """Simple inference to get output channel size."""
+        dummy_input = torch.zeros(1, 50, n_input, device=self._state.device)
+        with torch.no_grad():
+            output, _ = model(dummy_input)
+
+        if isinstance(output, dict):
+            return {k: v.shape[-1] for k, v in output.items()}
+        else:
+            return {"output": output.shape[-1]}
+
+    def _reset_state(self, message: AxisArray) -> None:
+        model_kwargs = dict(self.settings.model_kwargs or {})
+        self._common_reset_state(message, model_kwargs)
+        self._init_optimizer()
+        self._validate_loss_keys(list(self._state.chan_ax.keys()))
+
+        batch_size = 1 if message.data.ndim == 2 else message.data.shape[0]
+        self.reset_hidden(batch_size)
+
+    def _maybe_reset_state(self, message: AxisArray, batch_size: int) -> bool:
+        preserve_state = self.settings.preserve_state_across_windows
+        if preserve_state == "auto":
+            axes = message.axes
+            if batch_size < 2:
+                # Single window, so preserve
+                preserve_state = True
+            elif "time" not in axes or "win" not in axes:
+                # Default fallback
+                ez.logger.warning(
+                    "Missing 'time' or 'win' axis for auto preserve-state logic. Defaulting to reset."
+                )
+                preserve_state = False
+            else:
+                # Calculate stride between windows (assuming uniform spacing)
+                win_stride = axes["win"].gain
+                # Calculate window length from time axis length and gain
+                time_len = message.data.shape[message.get_axis_idx("time")]
+                gain = getattr(axes["time"], "gain", None)
+                if gain is None:
+                    ez.logger.warning(
+                        "Time axis gain not found, using default gain of 1.0."
+                    )
+                    gain = 1.0  # fallback default
+                win_len = time_len * gain
+                # Determine if we should preserve state
+                preserve_state = win_stride >= win_len
+
+        # Preserve if windows do NOT overlap: stride >= window length
+        if not preserve_state:
+            self.reset_hidden(batch_size)
+        else:
+            # If preserving state, only reset if batch size isn't 1
+            hx_batch_size = (
+                self._state.hx[0].shape[1]
+                if isinstance(self._state.hx, tuple)
+                else self._state.hx.shape[1]
+            )
+            if hx_batch_size != 1:
+                ez.logger.debug(
+                    f"Resetting hidden state due to batch size mismatch (hx: {hx_batch_size}, new: 1)"
+                )
+                self.reset_hidden(1)
+        return preserve_state
+
+    def _process(self, message: AxisArray) -> list[AxisArray]:
+        x = message.data
+        if not isinstance(x, torch.Tensor):
+            x = torch.tensor(
+                x,
+                dtype=torch.float32
+                if self.settings.single_precision
+                else torch.float64,
+                device=self._state.device,
+            )
+
+        # Add batch dimension if missing
+        x, added_batch_dim = self._ensure_batched(x)
+
+        batch_size = x.shape[0]
+        preserve_state = self._maybe_reset_state(message, batch_size)
+
+        with torch.no_grad():
+            # If we are preserving state and have multiple batches, process sequentially
+            if preserve_state and batch_size > 1:
+                y_data = {}
+                for x_batch in x:
+                    x_batch = x_batch.unsqueeze(0)
+                    y, self._state.hx = self._state.model(x_batch, hx=self._state.hx)
+                    for key, out in y.items():
+                        if key not in y_data:
+                            y_data[key] = []
+                        y_data[key].append(out.cpu().numpy())
+                # Concatenate outputs for each key
+                y_data = {
+                    key: np.concatenate(outputs, axis=0)
+                    for key, outputs in y_data.items()
+                }
+            else:
+                y, self._state.hx = self._state.model(x, hx=self._state.hx)
+                y_data = {
+                    key: (
+                        out.cpu().numpy().squeeze(0)
+                        if added_batch_dim
+                        else out.cpu().numpy()
+                    )
+                    for key, out in y.items()
+                }
+
+        return [
+            replace(
+                message,
+                data=out,
+                axes={**message.axes, "ch": self._state.chan_ax[key]},
+                key=key,
+            )
+            for key, out in y_data.items()
+        ]
+
+    def reset_hidden(self, batch_size: int) -> None:
+        self._state.hx = self._state.model.init_hidden(batch_size, self._state.device)
+
+    def _train_step(
+        self,
+        X: torch.Tensor,
+        y_targ: dict[str, torch.Tensor],
+        loss_fns: dict[str, torch.nn.Module],
+    ) -> None:
+        y_pred, self._state.hx = self._state.model(X, hx=self._state.hx)
+        if not isinstance(y_pred, dict):
+            y_pred = {"output": y_pred}
+
+        loss_weights = self.settings.loss_weights or {}
+        losses = []
+        for key in y_targ.keys():
+            loss_fn = loss_fns.get(key)
+            if loss_fn is None:
+                raise ValueError(f"Loss function for key '{key}' is not defined.")
+            if isinstance(loss_fn, torch.nn.CrossEntropyLoss):
+                loss = loss_fn(y_pred[key].permute(0, 2, 1), y_targ[key].long())
+            else:
+                loss = loss_fn(y_pred[key], y_targ[key])
+            weight = loss_weights.get(key, 1.0)
+            losses.append(loss * weight)
+
+        total_loss = sum(losses)
+        ez.logger.debug(
+            f"Training step loss: {total_loss.item()} (individual losses: {[loss.item() for loss in losses]})"
+        )
+
+        self._state.optimizer.zero_grad()
+        total_loss.backward()
+        self._state.optimizer.step()
+        if self._state.scheduler is not None:
+            self._state.scheduler.step()
+
+    def partial_fit(self, message: SampleMessage) -> None:
+        self._state.model.train()
+
+        X = self._to_tensor(message.sample.data)
+
+        # Add batch dimension if missing
+        X, batched = self._ensure_batched(X)
+
+        batch_size = X.shape[0]
+        preserve_state = self._maybe_reset_state(message.sample, batch_size)
+
+        y_targ = message.trigger.value
+        if not isinstance(y_targ, dict):
+            y_targ = {"output": y_targ}
+        y_targ = {key: self._to_tensor(value) for key, value in y_targ.items()}
+        # Add batch dimension to y_targ values if missing
+        if batched:
+            for key in y_targ:
+                y_targ[key] = y_targ[key].unsqueeze(0)
+
+        loss_fns = self.settings.loss_fn
+        if loss_fns is None:
+            raise ValueError("loss_fn must be provided in settings to use partial_fit")
+        if not isinstance(loss_fns, dict):
+            loss_fns = {k: loss_fns for k in y_targ.keys()}
+
+        with torch.set_grad_enabled(True):
+            if preserve_state:
+                self._train_step(X, y_targ, loss_fns)
+            else:
+                for i in range(batch_size):
+                    self._train_step(
+                        X[i].unsqueeze(0),
+                        {key: value[i].unsqueeze(0) for key, value in y_targ.items()},
+                        loss_fns,
+                    )
+
+        self._state.model.eval()
+        if self.settings.reset_hidden_on_fit:
+            self.reset_hidden(X.shape[0])
+
+
+class RNNUnit(
+    BaseAdaptiveTransformerUnit[
+        RNNSettings,
+        AxisArray,
+        AxisArray,
+        RNNProcessor,
+    ]
+):
+    SETTINGS = RNNSettings
+
+    @ez.subscriber(BaseAdaptiveTransformerUnit.INPUT_SIGNAL, zero_copy=True)
+    @ez.publisher(BaseAdaptiveTransformerUnit.OUTPUT_SIGNAL)
+    @profile_subpub(trace_oldest=False)
+    async def on_signal(self, message: AxisArray) -> typing.AsyncGenerator:
+        results = await self.processor.__acall__(message)
+        for result in results:
+            yield self.OUTPUT_SIGNAL, result

ezmsg/learn/process/sgd.py

@@ -0,0 +1,131 @@
+import typing
+
+import ezmsg.core as ez
+import numpy as np
+from ezmsg.sigproc.sampler import SampleMessage
+from ezmsg.sigproc.base import GenAxisArray
+from ezmsg.util.generator import consumer
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+from sklearn.exceptions import NotFittedError
+from sklearn.linear_model import SGDClassifier
+
+from ..util import ClassifierMessage
+
+
+@consumer
+def sgd_decoder(
+    alpha: float = 1.5e-5,
+    eta0: float = 1e-7,  # Lower than what you'd use for offline training.
+    loss: str = "squared_hinge",
+    label_weights: dict[str, float] | None = None,
+    settings_path: str | None = None,
+) -> typing.Generator[AxisArray | SampleMessage, ClassifierMessage | None, None]:
+    """
+    Passive Aggressive Classifier
+    Online Passive-Aggressive Algorithms <http://jmlr.csail.mit.edu/papers/volume7/crammer06a/crammer06a.pdf>
+    K. Crammer, O. Dekel, J. Keshat, S. Shalev-Shwartz, Y. Singer - JMLR (2006)
+
+    Args:
+        alpha: Maximum step size (regularization)
+        eta0: The initial learning rate for the 'adaptive’ schedules.
+        loss: The loss function to be used:
+            hinge: equivalent to PA-I in the reference paper.
+            squared_hinge: equivalent to PA-II in the reference paper.
+        label_weights: An optional dictionary of label names and their relative weight.
+            e.g., {'Go': 31.0, 'Stop': 0.5}
+            If this is None then settings_path must be provided and the pre-trained model
+        settings_path: Path to the stored sklearn model pkl file.
+
+    Returns:
+        Generator that accepts `SampleMessage` for incremental training (`partial_fit`) and yields None,
+        or `AxisArray` for inference (`predict`) and yields a `ClassifierMessage`.
+    """
+    # pre-init inputs and outputs
+    msg_out = ClassifierMessage(data=np.array([]), dims=[""])
+
+    # State variables:
+
+    if settings_path is not None:
+        import pickle
+
+        with open(settings_path, "rb") as f:
+            model = pickle.load(f)
+            if label_weights is not None:
+                model.class_weight = label_weights
+            # Overwrite eta0, probably with a value lower than what was used online.
+            model.eta0 = eta0
+    else:
+        model = SGDClassifier(
+            loss=loss,
+            alpha=alpha,
+            penalty="elasticnet",
+            learning_rate="adaptive",
+            eta0=eta0,
+            early_stopping=False,
+            class_weight=label_weights,
+        )
+
+    b_first_train = True
+    # TODO: template_out
+
+    while True:
+        msg_in: AxisArray | SampleMessage = yield msg_out
+
+        msg_out = None
+        if type(msg_in) is SampleMessage:
+            # SampleMessage used for training.
+            if not np.any(np.isnan(msg_in.sample.data)):
+                train_sample = msg_in.sample.data.reshape(1, -1)
+                if b_first_train:
+                    model.partial_fit(
+                        train_sample,
+                        [msg_in.trigger.value],
+                        classes=list(label_weights.keys()),
+                    )
+                    b_first_train = False
+                else:
+                    model.partial_fit(train_sample, [msg_in.trigger.value])
+        elif msg_in.data.size:
+            # AxisArray used for inference
+            if not np.any(np.isnan(msg_in.data)):
+                try:
+                    X = msg_in.data.reshape((msg_in.data.shape[0], -1))
+                    result = model._predict_proba_lr(X)
+                except NotFittedError:
+                    result = None
+                if result is not None:
+                    out_axes = {}
+                    if msg_in.dims[0] in msg_in.axes:
+                        out_axes[msg_in.dims[0]] = replace(
+                            msg_in.axes[msg_in.dims[0]],
+                            offset=msg_in.axes[msg_in.dims[0]].offset,
+                        )
+                    msg_out = ClassifierMessage(
+                        data=result,
+                        dims=msg_in.dims[:1] + ["labels"],
+                        axes=out_axes,
+                        labels=list(model.class_weight.keys()),
+                        key=msg_in.key,
+                    )
+
+
+class SGDDecoderSettings(ez.Settings):
+    alpha: float = 1e-5
+    eta0: float = 3e-4
+    loss: str = "hinge"
+    label_weights: dict[str, float] | None = None
+    settings_path: str | None = None
+
+
+class SGDDecoder(GenAxisArray):
+    SETTINGS = SGDDecoderSettings
+    INPUT_SAMPLE = ez.InputStream(SampleMessage)
+
+    # Method to be implemented by subclasses to construct the specific generator
+    def construct_generator(self):
+        self.STATE.gen = sgd_decoder(**self.SETTINGS.__dict__)
+
+    @ez.subscriber(INPUT_SAMPLE)
+    async def on_sample(self, msg: SampleMessage) -> None:
+        _ = self.STATE.gen.send(msg)

ezmsg/learn/process/sklearn.py

@@ -0,0 +1,274 @@
+import importlib
+import pickle
+import typing
+
+import ezmsg.core as ez
+import numpy as np
+import pandas as pd
+from ezmsg.sigproc.base import (
+    BaseAdaptiveTransformer,
+    BaseAdaptiveTransformerUnit,
+    processor_state,
+)
+from ezmsg.sigproc.sampler import SampleMessage
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
+
+class SklearnModelSettings(ez.Settings):
+    model_class: str
+    """
+    Full path to the sklearn model class
+    Example: 'sklearn.linear_model.LinearRegression'
+    """
+    model_kwargs: dict[str, typing.Any] = None
+    """
+    Additional keyword arguments to pass to the model constructor.
+    Example: {'fit_intercept': True, 'normalize': False}
+    """
+    checkpoint_path: str | None = None
+    """
+    Path to a checkpoint file to load the model from.
+    If provided, the model will be initialized from this checkpoint.
+    Example: 'path/to/checkpoint.pkl'
+    """
+    partial_fit_classes: np.ndarray | None = None
+    """
+    The full list of classes to use for partial_fit calls.
+    This must be provided to use `partial_fit` with classifiers.
+    """
+
+
+@processor_state
+class SklearnModelState:
+    model: typing.Any = None
+    chan_ax: AxisArray.CoordinateAxis | None = None
+
+
+class SklearnModelProcessor(
+    BaseAdaptiveTransformer[
+        SklearnModelSettings, AxisArray, AxisArray, SklearnModelState
+    ]
+):
+    """
+    Processor that wraps a scikit-learn, River, or HMMLearn model for use in the ezmsg framework.
+
+    This processor supports:
+    - `fit`, `partial_fit`, or River's `learn_many`/`learn_one` for training.
+    - `predict`, River's `predict_many`, or `predict_one` for inference.
+    - Optional model checkpoint loading and saving.
+
+    The processor expects and outputs `AxisArray` messages with a `"ch"` (channel) axis.
+
+    Settings:
+    ---------
+    model_class : str
+        Full path to the sklearn or River model class to use.
+        Example: "sklearn.linear_model.SGDClassifier" or "river.linear_model.LogisticRegression"
+
+    model_kwargs : dict[str, typing.Any], optional
+        Additional keyword arguments passed to the model constructor.
+
+    checkpoint_path : str, optional
+        Path to a pickle file to load a previously saved model. If provided, the model will
+        be restored from this path at startup.
+
+    partial_fit_classes : np.ndarray, optional
+        For classifiers that require all class labels to be specified during `partial_fit`.
+
+    Example:
+    -----------------------------
+    ```python
+    processor = SklearnModelProcessor(
+        settings=SklearnModelSettings(
+            model_class='sklearn.linear_model.SGDClassifier',
+            model_kwargs={'loss': 'log_loss'},
+            partial_fit_classes=np.array([0, 1]),
+        )
+    )
+    ```
+    """
+
+    def _init_model(self) -> None:
+        module_path, class_name = self.settings.model_class.rsplit(".", 1)
+        model_cls = getattr(importlib.import_module(module_path), class_name)
+        kwargs = self.settings.model_kwargs or {}
+        self._state.model = model_cls(**kwargs)
+
+    def save_checkpoint(self, path: str) -> None:
+        with open(path, "wb") as f:
+            pickle.dump(self._state.model, f)
+
+    def load_checkpoint(self, path: str) -> None:
+        try:
+            with open(path, "rb") as f:
+                self._state.model = pickle.load(f)
+        except Exception as e:
+            ez.logger.error(f"Failed to load model from {path}: {str(e)}")
+            raise RuntimeError(f"Failed to load model from {path}: {str(e)}") from e
+
+    def _reset_state(self, message: AxisArray) -> None:
+        # Try loading from checkpoint first
+        if self.settings.checkpoint_path:
+            self.load_checkpoint(self.settings.checkpoint_path)
+            n_input = message.data.shape[message.get_axis_idx("ch")]
+            if hasattr(self._state.model, "n_features_in_"):
+                expected = self._state.model.n_features_in_
+                if expected != n_input:
+                    raise ValueError(
+                        f"Model expects {expected} features, but got {n_input}"
+                    )
+        else:
+            # No checkpoint, initialize from scratch
+            self._init_model()
+
+    def partial_fit(self, message: SampleMessage) -> None:
+        X = message.sample.data
+        y = message.trigger.value
+        if self._state.model is None:
+            self._reset_state(message.sample)
+        if hasattr(self._state.model, "partial_fit"):
+            kwargs = {}
+            if self.settings.partial_fit_classes is not None:
+                kwargs["classes"] = self.settings.partial_fit_classes
+            self._state.model.partial_fit(X, y, **kwargs)
+        elif hasattr(self._state.model, "learn_many"):
+            df_X = pd.DataFrame(
+                {
+                    k: v
+                    for k, v in zip(
+                        message.sample.axes["ch"].data, message.sample.data.T
+                    )
+                }
+            )
+            name = (
+                message.trigger.value.axes["ch"].data[0]
+                if hasattr(message.trigger.value, "axes")
+                and "ch" in message.trigger.value.axes
+                else "target"
+            )
+            ser_y = pd.Series(
+                data=np.asarray(message.trigger.value.data).flatten(),
+                name=name,
+            )
+            self._state.model.learn_many(df_X, ser_y)
+        elif hasattr(self._state.model, "learn_one"):
+            # river's random forest does not support learn_many
+            for xi, yi in zip(X, y):
+                features = {f"f{i}": xi[i] for i in range(len(xi))}
+                self._state.model.learn_one(features, yi)
+        else:
+            raise NotImplementedError(
+                "Model does not support partial_fit or learn_many"
+            )
+
+    def fit(self, X: np.ndarray, y: np.ndarray) -> None:
+        if self._state.model is None:
+            dummy_msg = AxisArray(
+                data=X,
+                dims=["time", "ch"],
+                axes={
+                    "time": AxisArray.TimeAxis(fs=1.0),
+                    "ch": AxisArray.CoordinateAxis(
+                        data=np.array([f"ch_{i}" for i in range(X.shape[1])]),
+                        dims=["ch"],
+                    ),
+                },
+            )
+            self._reset_state(dummy_msg)
+        if hasattr(self._state.model, "fit"):
+            self._state.model.fit(X, y)
+        elif hasattr(self._state.model, "learn_many"):
+            df_X = pd.DataFrame(X, columns=[f"f{i}" for i in range(X.shape[1])])
+            ser_y = pd.Series(y.flatten(), name="target")
+            self._state.model.learn_many(df_X, ser_y)
+        elif hasattr(self._state.model, "learn_one"):
+            # river's random forest does not support learn_many
+            for xi, yi in zip(X, y):
+                features = {f"f{i}": xi[i] for i in range(len(xi))}
+                self._state.model.learn_one(features, yi)
+        else:
+            raise NotImplementedError("Model does not support fit or learn_many")
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        if self._state.model is None:
+            raise RuntimeError(
+                "Model has not been fit yet. Call `fit()` or `partial_fit()` before processing."
+            )
+        X = message.data
+        original_shape = X.shape
+        n_input = X.shape[message.get_axis_idx("ch")]
+
+        # Ensure X is 2D
+        X = X.reshape(-1, n_input)
+        if hasattr(self._state.model, "n_features_in_"):
+            expected = self._state.model.n_features_in_
+            if expected != n_input:
+                raise ValueError(
+                    f"Model expects {expected} features, but got {n_input}"
+                )
+
+        if hasattr(self._state.model, "predict"):
+            y_pred = self._state.model.predict(X)
+        elif hasattr(self._state.model, "predict_many"):
+            df_X = pd.DataFrame(X, columns=[f"f{i}" for i in range(X.shape[1])])
+            y_pred = self._state.model.predict_many(df_X)
+            y_pred = np.array(list(y_pred))
+        elif hasattr(self._state.model, "predict_one"):
+            # river's random forest does not support predict_many
+            y_pred = np.array(
+                [
+                    self._state.model.predict_one(
+                        {f"f{i}": xi[i] for i in range(len(xi))}
+                    )
+                    for xi in X
+                ]
+            )
+        else:
+            raise NotImplementedError("Model does not support predict or predict_many")
+
+        # For scalar outputs, ensure the output is 2D
+        if y_pred.ndim == 1:
+            y_pred = y_pred[:, np.newaxis]
+
+        output_shape = original_shape[:-1] + (y_pred.shape[-1],)
+        y_pred = y_pred.reshape(output_shape)
+
+        if self._state.chan_ax is None:
+            self._state.chan_ax = AxisArray.CoordinateAxis(
+                data=np.arange(output_shape[1]), dims=["ch"]
+            )
+
+        return replace(
+            message,
+            data=y_pred,
+            axes={**message.axes, "ch": self._state.chan_ax},
+        )
+
+
+class SklearnModelUnit(
+    BaseAdaptiveTransformerUnit[
+        SklearnModelSettings, AxisArray, AxisArray, SklearnModelProcessor
+    ]
+):
+    """
+    Unit wrapper for the `SklearnModelProcessor`.
+
+    This unit provides a plug-and-play interface for using a scikit-learn or River model
+    in an ezmsg graph-based system. It takes in `AxisArray` inputs and outputs predictions
+    in the same format, optionally performing training via `partial_fit` or `fit`.
+
+    Example:
+    --------
+    ```python
+    unit = SklearnModelUnit(
+        settings=SklearnModelSettings(
+            model_class='sklearn.linear_model.SGDClassifier',
+            model_kwargs={'loss': 'log_loss'},
+            partial_fit_classes=np.array([0, 1]),
+        )
+    )
+    ```
+    """
+
+    SETTINGS = SklearnModelSettings