Towards Unified Sequence Representation II: Move YAIB tasks and per-step labels to SequenceProcessorV2

docs/api/processors/pyhealth.processors.TimeseriesBinaryLabelProcessor.rst

@@ -1,7 +1,7 @@
 pyhealth.processors.TimeseriesBinaryLabelProcessor
 ==================================================
 
-Binary (0/1) label per time step (``"timeseries_binary"``).
+Binary (0/1) labels at time points (``"timeseries_binary"``).
 
 .. autoclass:: pyhealth.processors.TimeseriesBinaryLabelProcessor
     :members:

docs/api/processors/pyhealth.processors.TimeseriesLabelProcessor.rst

@@ -1,7 +1,7 @@
 pyhealth.processors.TimeseriesLabelProcessor
 ============================================
 
-Base processor for labels given per time step, returning ``(value, mask)`` on the time grid of the input time series.
+Base processor for labels given at time points: a restricted ``sequence_v2`` field with one ``series`` level of scalar values, packed per batch.
 
 .. autoclass:: pyhealth.processors.TimeseriesLabelProcessor
     :members:

docs/api/processors/pyhealth.processors.TimeseriesRegressionLabelProcessor.rst

@@ -1,7 +1,7 @@
 pyhealth.processors.TimeseriesRegressionLabelProcessor
 ======================================================
 
-Real-valued label per time step (``"timeseries_regression"``).
+Real-valued labels at time points (``"timeseries_regression"``).
 
 .. autoclass:: pyhealth.processors.TimeseriesRegressionLabelProcessor
     :members:

pyhealth/models/base_model.py

@@ -27,37 +27,28 @@
 )
 
 
-def _masked_step_loss(
+def _label_record_loss(
     loss_fn: Callable,
     logits: torch.Tensor,
-    target: tuple[torch.Tensor, torch.Tensor],
+    target: tuple[torch.Tensor, ...],
 ) -> torch.Tensor:
-    """Averages a per-step loss over the time steps that have a label.
+    """Averages a loss over the label records of a batch.
 
     Args:
-        loss_fn: Elementwise loss accepting ``reduction="none"``.
-        logits: Per-step logits of shape (batch, steps, 1).
-        target: Batched ``(value, mask)`` of shape (batch, label steps), where
-            padding and unlabeled steps have mask 0. Values at those steps are
-            ignored and may be NaN.
+        loss_fn: The loss, e.g. ``F.binary_cross_entropy_with_logits``.
+        logits: One prediction per label record, ``[ΣN, 1]``.
+        target: The packed batch of a
+            :class:`~pyhealth.processors.TimeseriesLabelProcessor`:
+            ``(values [ΣN, 1], time [ΣN], batch offsets [B + 1])``.
     """
-    # Temporary workaround: per-step labels arrive as a (value, mask) tuple so
-    # the existing single-tensor label API stays unchanged. A more unified
-    # handling of padded/masked label tensors across models is a follow-up.
-    value, mask = target
-    if value.shape[1] > logits.shape[1]:
+    values = target[0]
+    if logits.shape != values.shape:
         raise ValueError(
-            f"The label has {value.shape[1]} time steps but the input only "
-            f"{logits.shape[1]}; label and input must share the time grid."
+            f"Expected one prediction per label record, {tuple(values.shape)}, "
+            f"got {tuple(logits.shape)}. Per-step labels need a model that "
+            "predicts at each label's time."
         )
-    logits = logits[:, : value.shape[1], 0]
-    labeled = mask > 0
-    # Replace unlabeled values before the loss so a NaN there reaches neither
-    # the loss nor its gradient, then drop those steps from the sum.
-    value = torch.where(labeled, value, torch.zeros_like(value))
-    loss = loss_fn(logits, value, reduction="none")
-    loss = torch.where(labeled, loss, torch.zeros_like(loss))
-    return loss.sum() / labeled.sum().clamp(min=1)
+    return loss_fn(logits, values)
 
 
 class BaseModel(ABC, nn.Module):
@@ -193,13 +184,14 @@ def _resolve_mode(self, schema_entry: Any) -> str:
         if isinstance(schema_entry, str):
             return schema_entry.lower()
 
-        # Get class reference
+        # The most specific registered class: a subclass of a registered
+        # processor (e.g. the timeseries labels, which are SequenceProcessorV2s)
+        # resolves to its own name before its base's.
         cls = schema_entry if inspect.isclass(schema_entry) else schema_entry.__class__
-        for name, registered_cls in PROCESSOR_REGISTRY.items():
-            if cls is registered_cls or issubclass(
-                cls, registered_cls
-            ):  # allow subclassing
-                return name.lower()
+        for klass in cls.__mro__:
+            for name, registered_cls in PROCESSOR_REGISTRY.items():
+                if klass is registered_cls:
+                    return name.lower()
         raise ValueError(
             f"Cannot resolve mode from output_schema entry {schema_entry}. Use a supported string"
         )
@@ -313,16 +305,14 @@ def get_loss_function(self) -> Callable:
             - multiclass: `F.cross_entropy`
             - multilabel: `F.binary_cross_entropy_with_logits`
             - regression: `F.mse_loss`
-            - timeseries_binary: `F.binary_cross_entropy_with_logits` per
-              time step, averaged over the labeled steps
-            - timeseries_regression: `F.mse_loss` per time step, averaged
-              over the labeled steps
+            - timeseries_binary: `F.binary_cross_entropy_with_logits` over
+              the label records
+            - timeseries_regression: `F.mse_loss` over the label records
 
-        The timeseries losses take per-step logits of shape (batch, steps, 1)
-        and the batched ``(value, mask)`` label, each of shape (batch, steps),
-        produced by an output processor whose ``process()`` returns
-        ``(value, mask)`` per sample with mask 0 for steps without a label.
-        Values at masked steps are ignored and may be NaN.
+        The timeseries losses take one logit per label record, ``[ΣN, 1]``,
+        and the packed batch of a
+        :class:`~pyhealth.processors.TimeseriesLabelProcessor`, ``(values
+        [ΣN, 1], time [ΣN], batch offsets [B + 1])``.
 
         For binary and multilabel labels, a weight set with
         :meth:`set_pos_weight` is passed as ``pos_weight``.
@@ -353,9 +343,11 @@ def weighted_bce(input, target, **kwargs):
         elif mode == "regression":
             return F.mse_loss
         elif mode == "timeseries_binary":
-            return functools.partial(_masked_step_loss, F.binary_cross_entropy_with_logits)
+            return functools.partial(
+                _label_record_loss, F.binary_cross_entropy_with_logits
+            )
         elif mode == "timeseries_regression":
-            return functools.partial(_masked_step_loss, F.mse_loss)
+            return functools.partial(_label_record_loss, F.mse_loss)
         else:
             raise ValueError(f"Invalid mode: {mode}")
 
@@ -373,10 +365,10 @@ def prepare_y_prob(self, logits: torch.Tensor) -> torch.Tensor:
             - multilabel: a tensor of shape (batch_size, num_labels) with values
                 in [0, 1], which is obtained with `torch.sigmoid()`
             - regression: a tensor of shape (batch_size, 1) with raw logits
-            - timeseries_binary: a tensor of shape (batch_size, steps, 1)
-                with values in [0, 1] per time step, from `torch.sigmoid()`
-            - timeseries_regression: a tensor of shape (batch_size, steps, 1)
-                with raw logits per time step
+            - timeseries_binary: a tensor of shape (label records, 1) with
+                values in [0, 1], from `torch.sigmoid()`
+            - timeseries_regression: a tensor of shape (label records, 1)
+                with raw logits
 
         Args:
             logits (torch.Tensor): The predicted logit tensor.

pyhealth/models/rnn.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional, Tuple
+from typing import Optional
 
 import torch
 import torch.nn as nn
@@ -21,6 +21,9 @@
 
 from .embedding import EmbeddingModel
 
+# Label modes with one prediction per label record (TimeseriesLabelProcessor).
+_PER_STEP_MODES = ("timeseries_binary", "timeseries_regression")
+
 
 class RNNLayer(nn.Module):
     """Recurrent neural network layer.
@@ -86,7 +89,7 @@ def forward(
         self,
         x: torch.Tensor,
         mask: Optional[torch.Tensor] = None,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """Forward propagation.
 
         Args:
@@ -225,12 +228,9 @@ def __init__(
             raise ValueError("hidden_size is determined by hidden_dim")
         assert len(self.label_keys) == 1, "Only one label key is supported if RNN is initialized"
         self.label_key = self.label_keys[0]
-        # Per-step (timeseries) labels predict at every input step.
-        self._per_step = self.mode in ("timeseries_binary", "timeseries_regression")
-        if self._per_step and kwargs.get("bidirectional", False):
+        if self.mode in _PER_STEP_MODES:
             raise ValueError(
-                "A per-step label needs a causal model; bidirectional=True would "
-                "use future time steps."
+                "RNN does not support per-step (timeseries_*) labels."
             )
 
         self.embedding_model = EmbeddingModel(dataset, embedding_dim)
@@ -243,7 +243,7 @@ def __init__(
         output_size = self.get_output_size()
         self.fc = nn.Linear(len(self.feature_keys) * self.hidden_dim, output_size)
 
-    def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
+    def forward(self, **kwargs) -> dict[str, torch.Tensor]:
         """Forward propagation.
 
         The label `kwargs[self.label_key]` is a list of labels for each patient.
@@ -253,7 +253,7 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
                 all the feature keys and the label key.
 
         Returns:
-            Dict[str, torch.Tensor]: A dictionary with the following keys:
+            dict[str, torch.Tensor]: A dictionary with the following keys:
                 - loss: a scalar tensor representing the loss.
                 - y_prob: a tensor representing the predicted probabilities.
                 - y_true: a tensor representing the true labels.
@@ -309,30 +309,14 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
                 else:
                     mask = (torch.abs(x).sum(dim=-1) != 0).int()
 
-            outputs, x = self.rnn[feature_key](x, mask)
-            # Per-step labels use the output at every step, else the last one.
-            patient_emb.append(outputs if self._per_step else x)
-
-        if self._per_step:
-            # Features without a time axis (one step) are repeated over time.
-            steps = max(emb.shape[1] for emb in patient_emb)
-            for i, emb in enumerate(patient_emb):
-                if emb.shape[1] == 1:
-                    patient_emb[i] = emb.expand(-1, steps, -1)
-                elif emb.shape[1] != steps:
-                    raise ValueError(
-                        "Per-step labels need all time series features on the "
-                        "same time grid."
-                    )
-        patient_emb = torch.cat(patient_emb, dim=-1)
-        # (patient, label_size), or (patient, steps, label_size) per step
+            _, x = self.rnn[feature_key](x, mask)
+            patient_emb.append(x)
+
+        patient_emb = torch.cat(patient_emb, dim=1)
+        # (patient, label_size)
         logits = self.fc(patient_emb)
         # obtain y_true, loss, y_prob
-        if self._per_step:
-            # (value, mask) per step; the loss averages over labeled steps.
-            y_true = tuple(t.to(self.device) for t in kwargs[self.label_key])
-        else:
-            y_true = kwargs[self.label_key].to(self.device)
+        y_true = kwargs[self.label_key].to(self.device)
         loss = self.get_loss_function()(logits, y_true)
         y_prob = self.prepare_y_prob(logits)
         results = {"loss": loss, "y_prob": y_prob, "y_true": y_true, "logit": logits}
@@ -443,6 +427,10 @@ def __init__(
 
         assert len(self.label_keys) == 1, "Only one label key is supported"
         self.label_key = self.label_keys[0]
+        if self.mode in _PER_STEP_MODES:
+            raise ValueError(
+                "MultimodalRNN does not support per-step (timeseries_*) labels."
+            )
 
         self.embedding_model = EmbeddingModel(dataset, embedding_dim)
 
@@ -496,7 +484,7 @@ def _is_sequential_processor(self, processor) -> bool:
             TimeseriesProcessor,
         ))
 
-    def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
+    def forward(self, **kwargs) -> dict[str, torch.Tensor]:
         """Forward propagation handling mixed modalities.
 
         The label `kwargs[self.label_key]` is a list of labels for each patient.
@@ -506,7 +494,7 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
                 all the feature keys and the label key.
 
         Returns:
-            Dict[str, torch.Tensor]: A dictionary with the following keys:
+            dict[str, torch.Tensor]: A dictionary with the following keys:
                 - loss: a scalar tensor representing the loss.
                 - y_prob: a tensor representing the predicted probabilities.
                 - y_true: a tensor representing the true labels.