ezmsg-learn 1.0__py3-none-any.whl

ezmsg/learn/process/transformer.py

@@ -0,0 +1,222 @@
+import typing
+
+import ezmsg.core as ez
+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 TransformerSettings(TorchModelSettings):
+    model_class: str = "ezmsg.learn.model.transformer.TransformerModel"
+    """
+    Fully qualified class path of the model to be used.
+    This should be "ezmsg.learn.model.transformer.TransformerModel" for this.
+    """
+    autoregressive_head: str | None = None
+    """
+    The name of the output head used for autoregressive decoding.
+    This should match one of the keys in the model's output dictionary.
+    If None, the first output head will be used.
+    """
+    max_cache_len: int | None = 128
+    """
+    Maximum length of the target sequence cache for autoregressive decoding.
+    This limits the context length during decoding to prevent excessive memory usage.
+    If set to None, the cache will grow indefinitely.
+    """
+
+
+class TransformerState(TorchModelState):
+    ar_head: str | None = None
+    """
+    The name of the autoregressive head used for decoding.
+    This is set based on the `autoregressive_head` setting.
+    If None, the first output head will be used.
+    """
+    tgt_cache: typing.Optional[torch.Tensor] = None
+    """
+    Cache for the target sequence used in autoregressive decoding.
+    This is updated with each processed message to maintain context.
+    """
+
+
+class TransformerProcessor(
+    BaseAdaptiveTransformer[
+        TransformerSettings, AxisArray, AxisArray, TransformerState
+    ],
+    TorchProcessorMixin,
+    ModelInitMixin,
+):
+    @property
+    def has_decoder(self) -> bool:
+        return self.settings.model_kwargs.get("decoder_layers", 0) > 0
+
+    def reset_cache(self) -> None:
+        self._state.tgt_cache = None
+
+    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()))
+
+        self._state.tgt_cache = None
+        if (
+            self.settings.autoregressive_head is not None
+            and self.settings.autoregressive_head not in self._state.chan_ax
+        ):
+            raise ValueError(
+                f"Autoregressive head '{self.settings.autoregressive_head}' not found in target dictionary keys: {list(self._state.chan_ax.keys())}"
+            )
+        self._state.ar_head = (
+            self.settings.autoregressive_head
+            if self.settings.autoregressive_head is not None
+            else list(self._state.chan_ax.keys())[0]
+        )
+
+    def _process(self, message: AxisArray) -> list[AxisArray]:
+        # If has_decoder is False, fallback to regular processing
+        if not self.has_decoder:
+            return self._common_process(message)
+
+        x = self._to_tensor(message.data)
+        x, _ = self._ensure_batched(x)
+        if x.shape[0] > 1:
+            raise ValueError("Autoregressive decoding only supports batch size 1.")
+
+        with torch.no_grad():
+            y_pred = self._state.model(x, tgt=self._state.tgt_cache)
+
+        pred = y_pred[self._state.ar_head]
+        if self._state.tgt_cache is None:
+            self._state.tgt_cache = pred[:, -1:, :]
+        else:
+            self._state.tgt_cache = torch.cat(
+                [self._state.tgt_cache, pred[:, -1:, :]], dim=1
+            )
+        if self.settings.max_cache_len is not None:
+            if self._state.tgt_cache.shape[1] > self.settings.max_cache_len:
+                # Trim the cache to the maximum length
+                self._state.tgt_cache = self._state.tgt_cache[
+                    :, -self.settings.max_cache_len :, :
+                ]
+
+        if isinstance(y_pred, dict):
+            return [
+                replace(
+                    message,
+                    data=out.squeeze(0).cpu().numpy(),
+                    axes={**message.axes, "ch": self._state.chan_ax[key]},
+                    key=key,
+                )
+                for key, out in y_pred.items()
+            ]
+        else:
+            return [
+                replace(
+                    message,
+                    data=y_pred.squeeze(0).cpu().numpy(),
+                    axes={**message.axes, "ch": self._state.chan_ax["output"]},
+                )
+            ]
+
+    def partial_fit(self, message: SampleMessage) -> None:
+        self._state.model.train()
+
+        X = self._to_tensor(message.sample.data)
+        X, batched = self._ensure_batched(X)
+
+        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()}
+
+        weights = self.settings.loss_weights or {}
+
+        if self.has_decoder:
+            if X.shape[0] != 1:
+                raise ValueError("Autoregressive decoding only supports batch size 1.")
+
+            # Create shifted target for autoregressive head
+            tgt_tensor = y_targ[self._state.ar_head]
+            tgt = torch.cat(
+                [
+                    torch.zeros(
+                        (1, 1, tgt_tensor.shape[-1]),
+                        dtype=tgt_tensor.dtype,
+                        device=tgt_tensor.device,
+                    ),
+                    tgt_tensor[:, :-1, :],
+                ],
+                dim=1,
+            )
+
+            # Reset tgt_cache at start of partial_fit to avoid stale context
+            self.reset_cache()
+            y_pred = self._state.model(X, tgt=tgt)
+        else:
+            # For non-autoregressive models, use the model directly
+            y_pred = self._state.model(X)
+
+        if not isinstance(y_pred, dict):
+            y_pred = {"output": y_pred}
+
+        with torch.set_grad_enabled(True):
+            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 in settings."
+                    )
+                loss = loss_fn(y_pred[key], y_targ[key])
+                weight = weights.get(key, 1.0)
+                losses.append(loss * weight)
+            total_loss = sum(losses)
+
+            self._state.optimizer.zero_grad()
+            total_loss.backward()
+            self._state.optimizer.step()
+            if self._state.scheduler is not None:
+                self._state.scheduler.step()
+
+        self._state.model.eval()
+
+
+class TransformerUnit(
+    BaseAdaptiveTransformerUnit[
+        TransformerSettings,
+        AxisArray,
+        AxisArray,
+        TransformerProcessor,
+    ]
+):
+    SETTINGS = TransformerSettings
+
+    @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/util.py

@@ -0,0 +1,66 @@
+from enum import Enum
+from dataclasses import dataclass, field
+import typing
+
+from ezmsg.util.messages.axisarray import AxisArray
+import sklearn.linear_model
+import river.linear_model
+# from sklearn.neural_network import MLPClassifier
+
+
+class RegressorType(str, Enum):
+    ADAPTIVE = "adaptive"
+    STATIC = "static"
+
+
+class AdaptiveLinearRegressor(str, Enum):
+    LINEAR = "linear"
+    LOGISTIC = "logistic"
+    SGD = "sgd"
+    PAR = "par"  # passive-aggressive
+    # MLP = "mlp"
+
+
+class StaticLinearRegressor(str, Enum):
+    LINEAR = "linear"
+    RIDGE = "ridge"
+
+
+ADAPTIVE_REGRESSORS = {
+    AdaptiveLinearRegressor.LINEAR: river.linear_model.LinearRegression,
+    AdaptiveLinearRegressor.LOGISTIC: river.linear_model.LogisticRegression,
+    AdaptiveLinearRegressor.SGD: sklearn.linear_model.SGDRegressor,
+    AdaptiveLinearRegressor.PAR: sklearn.linear_model.PassiveAggressiveRegressor,
+    # AdaptiveLinearRegressor.MLP: MLPClassifier,
+}
+
+
+# Function to get a regressor by type and name
+def get_regressor(
+    regressor_type: typing.Union[RegressorType, str],
+    regressor_name: typing.Union[AdaptiveLinearRegressor, StaticLinearRegressor, str],
+):
+    if isinstance(regressor_type, str):
+        regressor_type = RegressorType(regressor_type)
+
+    if regressor_type == RegressorType.ADAPTIVE:
+        if isinstance(regressor_name, str):
+            regressor_name = AdaptiveLinearRegressor(regressor_name)
+        return ADAPTIVE_REGRESSORS[regressor_name]
+    elif regressor_type == RegressorType.STATIC:
+        if isinstance(regressor_name, str):
+            regressor_name = StaticLinearRegressor(regressor_name)
+        return STATIC_REGRESSORS[regressor_name]
+    else:
+        raise ValueError(f"Unknown regressor type: {regressor_type}")
+
+
+STATIC_REGRESSORS = {
+    StaticLinearRegressor.LINEAR: sklearn.linear_model.LinearRegression,
+    StaticLinearRegressor.RIDGE: sklearn.linear_model.Ridge,
+}
+
+
+@dataclass
+class ClassifierMessage(AxisArray):
+    labels: list[str] = field(default_factory=list)

ezmsg_learn-1.0.dist-info/METADATA

@@ -0,0 +1,34 @@
+Metadata-Version: 2.4
+Name: ezmsg-learn
+Version: 1.0
+Summary: ezmsg namespace package for machine learning
+Author-email: Chadwick Boulay <chadwick.boulay@gmail.com>
+License-Expression: MIT
+Requires-Python: >=3.10.15
+Requires-Dist: ezmsg-sigproc
+Requires-Dist: river>=0.22.0
+Requires-Dist: scikit-learn>=1.6.0
+Requires-Dist: torch>=2.6.0
+Description-Content-Type: text/markdown
+
+# ezmsg-learn
+
+This repository contains a Python package with modules for machine learning (ML)-related processing in the [`ezmsg`](https://www.ezmsg.org) framework. As ezmsg is intended primarily for processing unbounded streaming signals, so are the modules in this repo.
+
+> If you are only interested in offline analysis without concern for reproducibility in online applications, then you should probably look elsewhere.
+
+Processing units include dimensionality reduction, linear regression, and classification that can be initialized with known weights, or adapted on-the-fly with incoming (labeled) data. Machine-learning code depends on `river`, `scikit-learn`, `numpy`, and `torch`.
+
+## Getting Started
+
+This ezmsg namespace package is still highly experimental and under active development. It is not yet available on PyPI, so you will need to install it from source. The easiest way to do this is to use the `pip` command to install the package directly from GitHub:
+
+```bash
+pip install git+ssh://git@github.com/ezmsg-org/ezmsg-learn
+```
+
+Note that this package depends on a specific version of `ezmsg-sigproc` (specifically, [this branch]("70-use-protocols-for-axisarray-transformers")) that has yet to be merged and released. This may conflict with your project's separate dependency on ezmsg-sigproc. However, this specific version of ezmsg-sigproc should be backwards compatible with its main branch, so in your project you can modify the dependency on ezmsg-sigproc to point to the new branch. e.g.,
+
+```bash
+pip install git+ssh://git@github.com/ezmsg-org/ezmsg-sigproc@70-use-protocols-for-axisarray-transformers
+```

ezmsg_learn-1.0.dist-info/RECORD

@@ -0,0 +1,36 @@
+ezmsg/learn/__init__.py,sha256=9vTW4C2EQCHgDAo8gIlGNDfOxcQpNGV3Cct9-HsBJKY,57
+ezmsg/learn/__version__.py,sha256=T-YAefOAMONzdzJN9AfYa3q6PjJ-HRflYoFg45W1xFU,699
+ezmsg/learn/util.py,sha256=-WZ3k0sWSIJ1Z9aNiNFFYGa1-8oSbwQc-wI2i86w_C4,2014
+ezmsg/learn/dim_reduce/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ezmsg/learn/dim_reduce/adaptive_decomp.py,sha256=HmjantyGCIt7gP1d0iJ9mfoEtvR3-FjVrpv9JBT6mdc,9332
+ezmsg/learn/dim_reduce/incremental_decomp.py,sha256=FRx0Rhn3q8yHe64e4jaHSwWNVJ9eJT1ltuZWZc-C1R8,6830
+ezmsg/learn/linear_model/__init__.py,sha256=7_bcxc40W6UN2IfnJfjuVHe5mZ0BSPdHLKqCXpHPMwQ,78
+ezmsg/learn/linear_model/adaptive_linear_regressor.py,sha256=zfbwjTaBGdlBzzDBU6Nu2nuUjmhFVYs7Co84VWyRqIE,332
+ezmsg/learn/linear_model/cca.py,sha256=H-NnK9QH5vI0OnGezf674lV-X15xGzgar6d6G-67fZU,57
+ezmsg/learn/linear_model/linear_regressor.py,sha256=qTkIV2FRMC-Lhg7RTxmX1pWWsov84gV-47x-YSdNGdM,211
+ezmsg/learn/linear_model/sgd.py,sha256=6glInxmhapMVlBSZZB01w3vCqkNyGYOHoV8nfzrYCNI,138
+ezmsg/learn/linear_model/slda.py,sha256=MmmDbfm5y-8XpLThgJlHVV3f2kqZ-mrpnNkkIUHrwWg,151
+ezmsg/learn/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ezmsg/learn/model/cca.py,sha256=wf6vsfGHWiPt4FgZEG6veudKi1-GJDVp0bxaabxjJvE,4262
+ezmsg/learn/model/mlp.py,sha256=bMm4JX53sRdCLgMeh8gAvJnd2qruiTsyuv4MiwLpvCU,5574
+ezmsg/learn/model/mlp_old.py,sha256=ODpeoU-6DFG_yZLF9jdBOeRsw1ptyMyWCKM-WGJBeRs,2140
+ezmsg/learn/model/refit_kalman.py,sha256=PgjpMpIWkwHX7baSRLcT5loBA6h8uCf1L1spPfM4Nq8,16313
+ezmsg/learn/model/rnn.py,sha256=VmcRG2UcxROyrY10AJ-jwompOThtuj6SYjPzHWxAohw,6654
+ezmsg/learn/model/transformer.py,sha256=Vfo_CATOmoNtaqkE4BRWJy6QcX-1J_Pz3rMG2ErQSLM,7535
+ezmsg/learn/nlin_model/__init__.py,sha256=Ap62_lD7Dc3D-15ebhhBzT23PMrdyiV8fWV6eWy6wnE,101
+ezmsg/learn/nlin_model/mlp.py,sha256=zWTtI1JBP4KDcZCAqecHNn4Y59egEE9Bg2vi8-9la7k,165
+ezmsg/learn/process/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ezmsg/learn/process/adaptive_linear_regressor.py,sha256=0v6xRgWQWub0KlxOlb8ktYuohcwm4gvjHwOtv1qnNHg,5485
+ezmsg/learn/process/base.py,sha256=MhJstPfoTBNqLMDBAZqwbKueYyr9HJhvLzeU3k-Wl8E,7023
+ezmsg/learn/process/linear_regressor.py,sha256=2uKxvMraWITU0xZwbbI3O_mDD4cysl8RzHgl2XrPOPA,3119
+ezmsg/learn/process/mlp_old.py,sha256=b86ee-l0RxUSwCuTGe8MdAPbyJLOToIehAKKvf7bais,6958
+ezmsg/learn/process/refit_kalman.py,sha256=PWva0m33cPTjS8GX1aIK1tT5rgEHRub6Pp7GIgIhxsc,15220
+ezmsg/learn/process/rnn.py,sha256=tF3ejKqCResqBBWmSZoa18KpgB36KHhkYrTohWbs8w4,9862
+ezmsg/learn/process/sgd.py,sha256=DTg1MIA9M_K0qNTWF0oHhnMnn11jyL_MjvdAYKg8Gwk,4894
+ezmsg/learn/process/sklearn.py,sha256=5nyMoJffqpLin86h-tZicVIZUmKYxCt1xT2mr5vGkc4,9988
+ezmsg/learn/process/slda.py,sha256=M-zEiySPM4ovn0Os-ZaLRT8tyTTropsxoOj-veajSYg,4417
+ezmsg/learn/process/torch.py,sha256=JMHSoVbuRCpYdg7JuCOeauhQkDHmnjiZ-XfuHp2TZjw,13353
+ezmsg/learn/process/transformer.py,sha256=yhSDrREQy2PaiLVUzhwDppzg9g-NOox1gxatBMxj3BE,7872
+ezmsg_learn-1.0.dist-info/METADATA,sha256=1qQqXG-QgCb9yXuqg6gQi-re5mmPt6gh98RbwfJlJ94,2012
+ezmsg_learn-1.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+ezmsg_learn-1.0.dist-info/RECORD,,

ezmsg_learn-1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py3-none-any