Skip to content

Adds batching and CTCLoss functionality to RNNProcessor #16

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Sam-NT wants to merge 2 commits into
base: dev
Choose a base branch
from
Open

Adds batching and CTCLoss functionality to RNNProcessor #16

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
86ec9c2
RNNProcessor supports whole-batch training
Sam-NT Jul 7, 2026
f9ce2f8
RNNProcessor supports CTCLoss
Sam-NT Jul 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 32 additions & 2 deletions src/ezmsg/learn/process/rnn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,12 @@ class RNNSettings(TorchModelSettings):
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.
"""
batch_train: bool = False
"""
When True, train on the full batch in a single forward/backward pass
using packed sequences to handle varying sequence lengths.
When False, train each sample individually (current behavior).
"""
Comment on lines +38 to +43

@kylmcgr kylmcgr Jul 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Might be good to specify that when this is true, the partial fit message must contain data_len and trigger_len.



class RNNState(TorchModelState):
Expand Down Expand Up @@ -154,8 +160,10 @@ def _train_step(
X: torch.Tensor,
y_targ: dict[str, torch.Tensor],
loss_fns: dict[str, torch.nn.Module],
input_lens: torch.Tensor | None = None,
target_lens: torch.Tensor | None = None,
) -> None:
y_pred, self._state.hx = self._state.model(X, hx=self._state.hx)
y_pred, self._state.hx = self._state.model(X, hx=self._state.hx, input_lens=input_lens)

@kylmcgr kylmcgr Jul 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This will break existing custom RNN models that can be used with the RNN processor if they don't accept input_len argument in their forward pass. If we want to enforce that the models have the input_lens parameter, we can do that (I think its would just break a couple of my models, so I could update those), otherwise it might be good only pass the kwarg when it's actually set.

if not isinstance(y_pred, dict):
y_pred = {"output": y_pred}

Expand All @@ -167,6 +175,12 @@ def _train_step(
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())
elif isinstance(loss_fn, torch.nn.CTCLoss):
if input_lens is None or target_lens is None:
raise ValueError("CTCLoss requires input_lens and target_lens in batch training mode.")
log_probs = torch.nn.functional.log_softmax(y_pred[key], dim=-1).permute(1, 0, 2)
targets = y_targ[key].flatten().long()

@kylmcgr kylmcgr Jul 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Targets: Tensor of size (N,S) or (sum(target_lengths)), where N=batch size and S=max target length, if shape is (N,S). It represents the target sequences. Each element in the target sequence is a class index. And the target index cannot be blank (default=0). In the (N,S) form, targets are padded to the length of the longest sequence, and stacked. In the (sum(target_lengths)) form, the targets are assumed to be un-padded and concatenated within 1 dimension.

From the torch.nn.CTCLoss documentation, it looks like it expects since (N,S) or a sum of target_lengths. Correct me if I'm wrong, but this would lead me to believe that flattening the targets would not be the correct way to pass it into the loss function. Were you able to train a model like this that decoded well? If so, maybe I'm missing something about how you are passing the targets in.

loss = loss_fn(log_probs, targets, input_lengths=input_lens, target_lengths=target_lens)
else:
loss = loss_fn(y_pred[key], y_targ[key])
weight = loss_weights.get(key, 1.0)
Expand Down Expand Up @@ -210,7 +224,23 @@ def partial_fit(self, message: AxisArray) -> None:
loss_fns = {k: loss_fns for k in y_targ.keys()}

with torch.set_grad_enabled(True):
if preserve_state:
if self.settings.batch_train:
input_lens = message.attrs.get("data_len")
target_lens = message.attrs.get("trigger_len")
if input_lens is not None:
input_lens = torch.tensor(input_lens, dtype=torch.int64, device="cpu")
if target_lens is not None:
target_lens = torch.tensor(target_lens, dtype=torch.long, device=self._state.device)
self.reset_hidden(batch_size)
self._train_step(X, y_targ, loss_fns, input_lens=input_lens, target_lens=target_lens)
else:
ez.logger.warning(

@kylmcgr kylmcgr Jul 28, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If loss is CTC, then fall-back will error since it needs data_len. Consider if we should give an informative ezmsg logger error here rather than wait for it to hit the train step and error there.

"batch_train=True but 'data_len' not in message.attrs; falling back to per-sample training."
)
self.reset_hidden(batch_size)
for i in range(batch_size):
self._train_step(X[i].unsqueeze(0), {k: v[i].unsqueeze(0) for k, v in y_targ.items()}, loss_fns)
Comment on lines +241 to +242

@kylmcgr kylmcgr Jul 28, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This fall-back is for batch_train=True but no 'data_len'. Should this not still be a single train step for the batch rather than loop over the batch? Maybe this whole section can turn into something like this?

if self.settings.batch_train:
    input_lens = message.attrs.get("data_len")
    target_lens = message.attrs.get("trigger_len")

    if input_lens is None and any(isinstance(fn, torch.nn.CTCLoss) for fn in loss_fns.values()):
        raise ValueError(
            "CTCLoss requires 'data_len'/'trigger_len' in message.attrs."
        )
    if input_lens is None:
        ez.logger.warning(
            "batch_train=True but 'data_len' not in attrs; training on the full "
            "batch without length masking (padded timesteps contribute to the loss)."
        )

    if input_lens is not None:
        input_lens = torch.as_tensor(input_lens, dtype=torch.int64, device="cpu")
    if target_lens is not None:
        target_lens = torch.as_tensor(target_lens, dtype=torch.int64, device="cpu")

    self.reset_hidden(batch_size)
    self._train_step(X, y_targ, loss_fns, input_lens=input_lens, target_lens=target_lens)

elif preserve_state:
self._train_step(X, y_targ, loss_fns)
else:
for i in range(batch_size):
Expand Down
Loading