Wraps ConvSubsampling (8x downsampling) for per-chunk export. Input: chunk (1, T_mel_chunk, 128) float, chunk_len (1,) int64. Output: embs (1, T_sub, 512) float, emb_len (1,) int64. Re-implements ConvSubsampling.forward without MaskedConvSequential's masking logic, which creates da
| 44 | |
| 45 | |
| 46 | class PreEncodeWrapper(torch.nn.Module): |
| 47 | """Wraps ConvSubsampling (8x downsampling) for per-chunk export. |
| 48 | |
| 49 | Input: chunk (1, T_mel_chunk, 128) float, chunk_len (1,) int64. |
| 50 | Output: embs (1, T_sub, 512) float, emb_len (1,) int64. |
| 51 | |
| 52 | Re-implements ConvSubsampling.forward without MaskedConvSequential's |
| 53 | masking logic, which creates data-dependent guards that block torch.export. |
| 54 | Masking is unnecessary for single-sample inference with valid-length chunks. |
| 55 | """ |
| 56 | |
| 57 | def __init__(self, pre_encode): |
| 58 | super().__init__() |
| 59 | self.conv_layers = nn.ModuleList(list(pre_encode.conv)) |
| 60 | self.out = pre_encode.out |
| 61 | self._left_padding = pre_encode._left_padding |
| 62 | self._right_padding = pre_encode._right_padding |
| 63 | self._kernel_size = pre_encode._kernel_size |
| 64 | self._stride = pre_encode._stride |
| 65 | self._ceil_mode = pre_encode._ceil_mode |
| 66 | self._sampling_num = pre_encode._sampling_num |
| 67 | |
| 68 | def forward( |
| 69 | self, chunk: torch.Tensor, chunk_len: torch.Tensor |
| 70 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 71 | from nemo.collections.asr.parts.submodules.subsampling import calc_length |
| 72 | |
| 73 | out_lengths = calc_length( |
| 74 | chunk_len, |
| 75 | all_paddings=self._left_padding + self._right_padding, |
| 76 | kernel_size=self._kernel_size, |
| 77 | stride=self._stride, |
| 78 | ceil_mode=self._ceil_mode, |
| 79 | repeat_num=self._sampling_num, |
| 80 | ) |
| 81 | x = chunk.unsqueeze(1) # (B, 1, T, feat_in) |
| 82 | for layer in self.conv_layers: |
| 83 | x = layer(x) |
| 84 | # flatten static dims C and F (not dynamic T) to avoid symbolic guards |
| 85 | x = x.permute(0, 2, 1, 3).flatten(2) # (B, T_sub, C*F) |
| 86 | x = self.out(x) |
| 87 | return x, out_lengths.to(torch.int64) |
| 88 | |
| 89 | |
| 90 | class EncodeWrapper(torch.nn.Module): |