From 738768e608a6b2b2d53e671187f33531b39e76c0 Mon Sep 17 00:00:00 2001 From: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com> Date: Mon, 7 Sep 2026 12:32:47 +0530 Subject: [PATCH] fix(asr): repair ConvSubsampling forward paths missed by the MaskedConvSequential refactor The MaskedConvSequential refactor left three ConvSubsampling paths broken: 1. `subsampling_conv_chunking_factor=-1` (chunking disabled, documented in the class docstring) falls into `x, lengths = self.conv(x)`, which calls MaskedConvSequential.forward without the required `lengths` argument and raises TypeError for every conv2d stack (vggnet, striding, dw_striding). 2. The 1-D stacks (striding_conv1d, dw_striding_conv1d, conv2d_subsampling=False) always take that same branch regardless of the chunking factor, so both variants raise TypeError on every forward pass. 3. For vggnet, `_forward_torch` reads `layer.kernel_size[0]` / `layer.stride[0]`, but nn.MaxPool2d stores the int values passed in (kernel_size=2, stride=2), so every vggnet forward raises TypeError: 'int' object is not subscriptable. _layer_padding has the same problem with the int `padding`. Once the int access works, the length update must also honor MaxPool2d's ceil_mode=True, which calculate_conv_output_size's floor division ignores. Fixes: - ConvSubsampling.forward: pass `lengths` to the masked stack when chunking is disabled; run 1-D stacks as a plain sequential (MaskedConvSequential.forward accepts lengths=None for exactly this) and report the precomputed out_lengths. - _forward_torch: read kernel/stride/padding via a _pair_first helper that tolerates int attributes, and add the ceil-mode remainder frame when a pooling layer runs with ceil_mode=True. Adds CPU unit tests for all three paths; each fails with TypeError or wrong lengths on the unfixed code. The chunking=-1 path is asserted to match the default (chunked) path bit-for-bit, and conv1d/vggnet lengths are asserted against calc_length references. Signed-off-by: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com> --- .../asr/parts/submodules/subsampling.py | 55 +++++++-- tests/collections/asr/test_asr_subsampling.py | 104 ++++++++++++++++++ 2 files changed, 150 insertions(+), 9 deletions(-) diff --git a/nemo/collections/asr/parts/submodules/subsampling.py b/nemo/collections/asr/parts/submodules/subsampling.py index 3ce6950b981e..17bfc36527b2 100644 --- a/nemo/collections/asr/parts/submodules/subsampling.py +++ b/nemo/collections/asr/parts/submodules/subsampling.py @@ -492,7 +492,13 @@ def forward(self, x, lengths): else: x, lengths = self.conv(x, lengths) else: - x, lengths = self.conv(x) + if self.conv2d_subsampling: + # Chunking disabled (-1): run the masked conv stack directly. + x, lengths = self.conv(x, lengths) + else: + # 1-D conv stacks run without masking; keep the lengths computed above. + x = self.conv(x) + lengths = out_lengths # Flatten Channel and Frequency Axes if self.conv2d_subsampling: @@ -725,7 +731,15 @@ class MaskedConvSequential(nn.Sequential): # Set by ConvSubsampling; off by default, so every other subsampling type stays on PyTorch. fuse_triton = False - def forward(self, x, lengths): + def forward(self, x, lengths=None): + if lengths is None: + # Plain pass-through, equivalent to nn.Sequential. Used by the 1-D conv + # stacks (striding_conv1d / dw_striding_conv1d), whose inputs are already + # channels-first and which do not take part in length masking. + for layer in self: + x = layer(x) + return x + # Convert input (batch, time, features) to conv format x = x.unsqueeze(1) # (batch, 1, time, features) current_lengths = lengths @@ -760,10 +774,22 @@ def _forward_torch(self, x, current_lengths): x = layer(x) # Update lengths for stride operations with proper padding - if hasattr(layer, 'stride') and layer.stride != (1, 1): - current_lengths = calculate_conv_output_size( - current_lengths, layer.kernel_size[0], layer.stride[0], _layer_padding(layer) - ) + if hasattr(layer, 'stride') and _pair_first(layer.stride) != 1: + kernel_size = _pair_first(layer.kernel_size) + stride = _pair_first(layer.stride) + left_pad, right_pad = _layer_padding(layer) + if getattr(layer, 'ceil_mode', False): + # Ceil-mode pooling (e.g. vggnet's MaxPool2d) emits one extra frame + # whenever the floor division drops a remainder. + remainder = (current_lengths + left_pad + right_pad - kernel_size) % stride + current_lengths = calculate_conv_output_size( + current_lengths, kernel_size, stride, (left_pad, right_pad) + ) + current_lengths = current_lengths + (remainder != 0).long() + else: + current_lengths = calculate_conv_output_size( + current_lengths, kernel_size, stride, (left_pad, right_pad) + ) mask = self._create_mask(x, current_lengths.long()) return x, current_lengths, mask @@ -826,15 +852,26 @@ def _create_mask(self, tensor, lengths): return time_mask.unsqueeze(-1).expand(batch_size, time, features).to(tensor.dtype) +def _pair_first(value): + """First element of an int-or-tuple kernel/stride/padding attribute. + + nn.Conv2d stores tuples, but the pooling modules (nn.MaxPool2d/nn.AvgPool2d) keep + whatever was passed in, which ConvSubsampling passes as a plain int. + """ + return value[0] if isinstance(value, tuple) else value + + def _layer_padding(layer): """The (start, end) padding of a convolution. - nn.Conv2d's `.padding` is (pad_h, pad_w), one value per axis and symmetric within it, so the - height value is both edges. CausalConv2D keeps its two edges on private attributes. + nn.Conv2d's `.padding` is (pad_h, pad_w), one value per axis and symmetric within it, so + the height value is both edges. CausalConv2D keeps its two edges on private attributes. + Pooling layers store an int (or an int-per-axis tuple), which is symmetric too. """ if hasattr(layer, "_left_padding"): return layer._left_padding, layer._right_padding - return layer.padding[0], layer.padding[0] + padding = _pair_first(layer.padding) + return padding, padding def _is_depthwise(layer): diff --git a/tests/collections/asr/test_asr_subsampling.py b/tests/collections/asr/test_asr_subsampling.py index b9fa1ee29235..6fa968a06a60 100644 --- a/tests/collections/asr/test_asr_subsampling.py +++ b/tests/collections/asr/test_asr_subsampling.py @@ -16,6 +16,7 @@ import torch from nemo.collections.asr.models import ASRModel +from nemo.collections.asr.parts.submodules.subsampling import ConvSubsampling, calc_length class TestASRSubsamplingConvChunking: @@ -60,3 +61,106 @@ def test_forward(self): assert diff <= 0.2 diff = torch.mean(torch.abs(logprobs_batch4_split - logprobs_batch4_nosplit)) assert diff <= 0.2 + + +class TestConvSubsamplingForwardPaths: + """CPU tests for ConvSubsampling paths that the chunking/splitting tests do not cover. + + Covers `subsampling_conv_chunking_factor=-1` (chunking disabled), the 1-D conv + stacks, and ceil-mode length bookkeeping for vggnet pooling. + """ + + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + @pytest.mark.parametrize("subsampling", ["striding", "dw_striding", "vggnet"]) + @pytest.mark.parametrize("subsampling_factor", [4, 8]) + def test_no_chunking_forward_matches_chunked_path(self, subsampling, subsampling_factor): + """With chunking disabled (-1), forward must run and match the default path.""" + torch.manual_seed(0) + no_chunk = ConvSubsampling( + subsampling=subsampling, + subsampling_factor=subsampling_factor, + feat_in=80, + feat_out=64, + conv_channels=32, + subsampling_conv_chunking_factor=-1, + ).eval() + default = ConvSubsampling( + subsampling=subsampling, + subsampling_factor=subsampling_factor, + feat_in=80, + feat_out=64, + conv_channels=32, + ).eval() + default.load_state_dict(no_chunk.state_dict()) + + x = torch.randn(2, 101, 80) + lengths = torch.tensor([101, 97]) + + with torch.inference_mode(): + out, out_lengths = no_chunk(x, lengths) + ref_out, ref_lengths = default(x, lengths) + + assert out.shape == ref_out.shape + assert torch.allclose(out, ref_out) + assert out_lengths.tolist() == ref_lengths.tolist() + assert bool((out_lengths <= out.shape[1]).all()) + + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + @pytest.mark.parametrize("subsampling", ["striding_conv1d", "dw_striding_conv1d"]) + @pytest.mark.parametrize("subsampling_factor", [2, 4, 8]) + def test_conv1d_stacks_forward(self, subsampling, subsampling_factor): + """The 1-D conv stacks must run with the default config and report calc_length lengths.""" + module = ConvSubsampling( + subsampling=subsampling, + subsampling_factor=subsampling_factor, + feat_in=80, + feat_out=64, + conv_channels=32, + ).eval() + + x = torch.randn(2, 101, 80) + lengths = torch.tensor([101, 97]) + + with torch.inference_mode(): + out, out_lengths = module(x, lengths) + + sampling_num = module._sampling_num + expected = calc_length( + lengths=lengths.to(dtype=torch.float), + all_paddings=module._left_padding + module._right_padding, + kernel_size=module._kernel_size, + stride=module._stride, + ceil_mode=module._ceil_mode, + repeat_num=sampling_num, + ) + assert out.shape[0] == 2 + assert out.shape[2] == 64 + assert out_lengths.tolist() == expected.tolist() + + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + @pytest.mark.parametrize("input_length", [101, 97]) + def test_vggnet_lengths_use_ceil_mode(self, input_length): + """vggnet pooling runs with ceil_mode=True; reported lengths must match it.""" + module = ConvSubsampling( + subsampling='vggnet', subsampling_factor=4, feat_in=80, feat_out=64, conv_channels=32 + ).eval() + + x = torch.randn(1, input_length, 80) + lengths = torch.tensor([input_length]) + + with torch.inference_mode(): + out, out_lengths = module(x, lengths) + + expected = calc_length( + lengths=lengths.to(dtype=torch.float), + all_paddings=0, + kernel_size=2, + stride=2, + ceil_mode=True, + repeat_num=2, + ) + assert out.shape[1] == expected.item() + assert out_lengths.tolist() == expected.tolist()