Adds batching and CTCLoss functionality to RNNProcessor - #16
Conversation
kylmcgr
left a comment
kylmcgr
left a comment
There was a problem hiding this comment.
Could you also add some test coverage for new functionality you added? Thanks!
| 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). | ||
| """ |
There was a problem hiding this comment.
Might be good to specify that when this is true, the partial fit message must contain data_len and trigger_len.
| 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) |
There was a problem hiding this comment.
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 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() |
There was a problem hiding this comment.
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.
| 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( |
There was a problem hiding this comment.
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.
| 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) |
There was a problem hiding this comment.
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)
2 participants
This PR adds support for
RNNProcessor.partial_fitto fit on a batch (multiple samples of sequences of features) instead of one sample at a time. It further adds support for usingtorch.nn.CTCLossas a customized loss function (similar totorch.nn.CrossEntropyLoss), due to its atypical function arguments.I roughly designed the data structure containing sequence and target lengths to match the existing method for passing targets and data in one AxisArray, but it's not the cleanest. Happy to work through a better implementation in this PR if desired.
Also worth noting that
torch.nn.CTCLossis not implemented in MPS in the latest release, but is in active development. The latest nightly worked for me.