ezmsg-learn 1.1.0__py3-none-any.whl

ezmsg/learn/process/rnn.py

@@ -0,0 +1,245 @@
+import typing
+
+import ezmsg.core as ez
+import numpy as np
+import torch
+from ezmsg.baseproc import BaseAdaptiveTransformer, BaseAdaptiveTransformerUnit
+from ezmsg.baseproc.util.profile import profile_subpub
+from ezmsg.sigproc.sampler import SampleMessage
+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,117 @@
+import typing
+
+import ezmsg.core as ez
+import numpy as np
+from ezmsg.baseproc import (
+    BaseAdaptiveTransformer,
+    BaseAdaptiveTransformerUnit,
+    SampleMessage,
+    processor_state,
+)
+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
+
+
+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
+
+
+@processor_state
+class SGDDecoderState:
+    model: typing.Any = None
+    b_first_train: bool = True
+
+
+class SGDDecoderTransformer(BaseAdaptiveTransformer[SGDDecoderSettings, AxisArray, ClassifierMessage, SGDDecoderState]):
+    """
+    SGD-based online 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)
+    """
+
+    def _refreshed_model(self):
+        if self.settings.settings_path is not None:
+            import pickle
+
+            with open(self.settings.settings_path, "rb") as f:
+                model = pickle.load(f)
+                if self.settings.label_weights is not None:
+                    model.class_weight = self.settings.label_weights
+                model.eta0 = self.settings.eta0
+        else:
+            model = SGDClassifier(
+                loss=self.settings.loss,
+                alpha=self.settings.alpha,
+                penalty="elasticnet",
+                learning_rate="adaptive",
+                eta0=self.settings.eta0,
+                early_stopping=False,
+                class_weight=self.settings.label_weights,
+            )
+        return model
+
+    def _reset_state(self, message: AxisArray) -> None:
+        self._state.model = self._refreshed_model()
+
+    def _process(self, message: AxisArray) -> ClassifierMessage | None:
+        if self._state.model is None or not message.data.size:
+            return None
+        if np.any(np.isnan(message.data)):
+            return None
+        try:
+            X = message.data.reshape((message.data.shape[0], -1))
+            result = self._state.model._predict_proba_lr(X)
+        except NotFittedError:
+            return None
+        out_axes = {}
+        if message.dims[0] in message.axes:
+            out_axes[message.dims[0]] = replace(
+                message.axes[message.dims[0]],
+                offset=message.axes[message.dims[0]].offset,
+            )
+        return ClassifierMessage(
+            data=result,
+            dims=message.dims[:1] + ["labels"],
+            axes=out_axes,
+            labels=list(self._state.model.class_weight.keys()),
+            key=message.key,
+        )
+
+    def partial_fit(self, message: SampleMessage) -> None:
+        if self._hash != 0:
+            self._reset_state(message.sample)
+            self._hash = 0
+
+        if np.any(np.isnan(message.sample.data)):
+            return
+        train_sample = message.sample.data.reshape(1, -1)
+        if self._state.b_first_train:
+            self._state.model.partial_fit(
+                train_sample,
+                [message.trigger.value],
+                classes=list(self.settings.label_weights.keys()),
+            )
+            self._state.b_first_train = False
+        else:
+            self._state.model.partial_fit(train_sample, [message.trigger.value])
+
+
+class SGDDecoder(
+    BaseAdaptiveTransformerUnit[
+        SGDDecoderSettings,
+        AxisArray,
+        ClassifierMessage,
+        SGDDecoderTransformer,
+    ]
+):
+    SETTINGS = SGDDecoderSettings

ezmsg/learn/process/sklearn.py

@@ -0,0 +1,241 @@
+import importlib
+import pickle
+import typing
+
+import ezmsg.core as ez
+import numpy as np
+import pandas as pd
+from ezmsg.baseproc 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