Files
wehub-resource-sync ba4be087d5
Create PR to main with cherry-pick from release / cherry-pick (push) Failing after 0s
CICD NeMo / pre-flight (push) Failing after 0s
CICD NeMo / configure (push) Has been skipped
Build, validate, and release Neural Modules / pre-flight (push) Failing after 1s
CICD NeMo / code-linting (push) Has been skipped
Build, validate, and release Neural Modules / release (push) Has been skipped
Build, validate, and release Neural Modules / release-summary (push) Has been cancelled
CICD NeMo / cicd-test-container-build (push) Has been cancelled
CICD NeMo / cicd-import-tests (push) Has been cancelled
CICD NeMo / L0_Setup_Test_Data_And_Models (push) Has been cancelled
CICD NeMo / cicd-main-unit-tests (push) Has been cancelled
CICD NeMo / cicd-main-speech (push) Has been cancelled
CICD NeMo / Nemo_CICD_Test (push) Has been cancelled
CICD NeMo / Coverage (e2e) (push) Has been cancelled
CICD NeMo / Coverage (unit-test) (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled
CICD NeMo / cicd-wait-in-queue (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 13:28:58 +08:00

399 lines
14 KiB
Python

# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import pytest
import torch
from nemo.collections.asr.parts.submodules import jasper
class TestJasperBlock:
@staticmethod
def jasper_base_config(**kwargs):
base = dict(
inplanes=16,
planes=8,
kernel_size=[11],
repeat=1,
stride=[1],
dilation=[1],
activation="relu",
conv_mask=True,
separable=False,
se=False,
)
base.update(kwargs)
return base
def check_module_exists(self, module, cls):
global _MODULE_EXISTS
_MODULE_EXISTS = 0
def _traverse(m):
if isinstance(m, cls):
global _MODULE_EXISTS
_MODULE_EXISTS += 1
module.apply(_traverse)
assert _MODULE_EXISTS > 0
@pytest.mark.unit
def test_basic_block(self):
config = self.jasper_base_config(residual=False)
act = jasper.jasper_activations.get(config.pop('activation'))()
block = jasper.JasperBlock(**config, activation=act)
x = torch.randn(1, 16, 131)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
assert isinstance(block, jasper.JasperBlock)
assert y[0].shape == torch.Size([1, config['planes'], 131])
assert ylen[0] == 131
@pytest.mark.unit
def test_residual_block(self):
config = self.jasper_base_config(residual=True)
act = jasper.jasper_activations.get(config.pop('activation'))()
block = jasper.JasperBlock(**config, activation=act)
x = torch.randn(1, 16, 131)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
assert isinstance(block, jasper.JasperBlock)
assert y[0].shape == torch.Size([1, config['planes'], 131])
assert ylen[0] == 131
@pytest.mark.unit
def test_basic_block_repeat(self):
config = self.jasper_base_config(residual=False, repeat=3)
act = jasper.jasper_activations.get(config.pop('activation'))()
block = jasper.JasperBlock(**config, activation=act)
x = torch.randn(1, 16, 131)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
assert isinstance(block, jasper.JasperBlock)
assert y[0].shape == torch.Size([1, config['planes'], 131])
assert ylen[0] == 131
assert len(block.mconv) == 3 * 3 + 1 # (3 repeats x {1 conv + 1 norm + 1 dropout} + final conv)
@pytest.mark.unit
def test_basic_block_repeat_stride(self):
config = self.jasper_base_config(residual=False, repeat=3, stride=[2])
act = jasper.jasper_activations.get(config.pop('activation'))()
block = jasper.JasperBlock(**config, activation=act)
x = torch.randn(1, 16, 131)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
assert isinstance(block, jasper.JasperBlock)
assert y[0].shape == torch.Size([1, config['planes'], 17]) # 131 // (stride ^ repeats)
assert ylen[0] == 17 # 131 // (stride ^ repeats)
assert len(block.mconv) == 3 * 3 + 1 # (3 repeats x {1 conv + 1 norm + 1 dropout} + final conv)
@pytest.mark.unit
def test_basic_block_repeat_stride_last(self):
config = self.jasper_base_config(residual=False, repeat=3, stride=[2], stride_last=True)
act = jasper.jasper_activations.get(config.pop('activation'))()
block = jasper.JasperBlock(**config, activation=act)
x = torch.randn(1, 16, 131)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
assert isinstance(block, jasper.JasperBlock)
assert y[0].shape == torch.Size([1, config['planes'], 66]) # 131 // stride
assert ylen[0] == 66 # 131 // stride
assert len(block.mconv) == 3 * 3 + 1 # (3 repeats x {1 conv + 1 norm + 1 dropout} + final conv)
@pytest.mark.unit
def test_basic_block_repeat_separable(self):
config = self.jasper_base_config(residual=False, repeat=3, separable=True)
act = jasper.jasper_activations.get(config.pop('activation'))()
block = jasper.JasperBlock(**config, activation=act)
x = torch.randn(1, 16, 131)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
assert isinstance(block, jasper.JasperBlock)
assert y[0].shape == torch.Size([1, config['planes'], 131])
assert ylen[0] == 131
assert len(block.mconv) == 3 * 4 + 1 # (3 repeats x {1 dconv + 1 pconv + 1 norm + 1 dropout} + final conv)
@pytest.mark.unit
def test_basic_block_stride(self):
config = self.jasper_base_config(stride=[2], residual=False)
act = jasper.jasper_activations.get(config.pop('activation'))()
print(config)
block = jasper.JasperBlock(**config, activation=act)
x = torch.randn(1, 16, 131)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
assert isinstance(block, jasper.JasperBlock)
assert y[0].shape == torch.Size([1, config['planes'], 66])
assert ylen[0] == 66
@pytest.mark.unit
def test_residual_block_stride(self):
config = self.jasper_base_config(stride=[2], residual=True, residual_mode='stride_add')
act = jasper.jasper_activations.get(config.pop('activation'))()
print(config)
block = jasper.JasperBlock(**config, activation=act)
x = torch.randn(1, 16, 131)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
assert isinstance(block, jasper.JasperBlock)
assert y[0].shape == torch.Size([1, config['planes'], 66])
assert ylen[0] == 66
@pytest.mark.unit
def test_residual_block_activations(self):
for activation in jasper.jasper_activations.keys():
config = self.jasper_base_config(activation=activation)
act = jasper.jasper_activations.get(config.pop('activation'))()
block = jasper.JasperBlock(**config, activation=act)
x = torch.randn(1, 16, 131)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
self.check_module_exists(block, act.__class__)
assert isinstance(block, jasper.JasperBlock)
assert y[0].shape == torch.Size([1, config['planes'], 131])
assert ylen[0] == 131
@pytest.mark.unit
def test_residual_block_normalizations(self):
NORMALIZATIONS = ["batch", "layer", "group"]
for normalization in NORMALIZATIONS:
config = self.jasper_base_config(normalization=normalization)
act = jasper.jasper_activations.get(config.pop('activation'))()
block = jasper.JasperBlock(**config, activation=act)
x = torch.randn(1, 16, 131)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
assert isinstance(block, jasper.JasperBlock)
assert y[0].shape == torch.Size([1, config['planes'], 131])
assert ylen[0] == 131
@pytest.mark.unit
def test_residual_block_se(self):
config = self.jasper_base_config(se=True, se_reduction_ratio=8)
act = jasper.jasper_activations.get(config.pop('activation'))()
block = jasper.JasperBlock(**config, activation=act)
x = torch.randn(1, 16, 131)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
self.check_module_exists(block, jasper.SqueezeExcite)
assert isinstance(block, jasper.JasperBlock)
assert y[0].shape == torch.Size([1, config['planes'], 131])
assert ylen[0] == 131
@pytest.mark.unit
def test_residual_block_asymmetric_pad_future_contexts(self):
# test future contexts at various values
# 0 = no future context
# 2 = limited future context
# 5 = symmetric context
# 8 = excess future context (more future context than present or past context)
future_contexts = [0, 2, 5, 8]
for future_context in future_contexts:
print(future_context)
config = self.jasper_base_config(future_context=future_context)
act = jasper.jasper_activations.get(config.pop('activation'))()
block = jasper.JasperBlock(**config, activation=act)
x = torch.randn(1, 16, 131)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
self.check_module_exists(block, torch.nn.ConstantPad1d)
self.check_module_exists(block, jasper.MaskedConv1d)
assert isinstance(block, jasper.JasperBlock)
assert y[0].shape == torch.Size([1, config['planes'], 131])
assert ylen[0] == 131
assert block.mconv[0].pad_layer is not None
assert block.mconv[0]._padding == (config['kernel_size'][0] - 1 - future_context, future_context)
@pytest.mark.unit
def test_residual_block_asymmetric_pad_future_context_fallback(self):
# test future contexts at various values
# 15 = K < FC; fall back to symmetric context
future_context = 15
print(future_context)
config = self.jasper_base_config(future_context=future_context)
act = jasper.jasper_activations.get(config.pop('activation'))()
block = jasper.JasperBlock(**config, activation=act)
x = torch.randn(1, 16, 131)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
self.check_module_exists(block, jasper.MaskedConv1d)
assert isinstance(block, jasper.JasperBlock)
assert y[0].shape == torch.Size([1, config['planes'], 131])
assert ylen[0] == 131
assert block.mconv[0].pad_layer is None
assert block.mconv[0]._padding == config['kernel_size'][0] // 2
@pytest.mark.unit
def test_padding_size_conv1d(self):
input_channels = 1
output_channels = 1
kernel_sizes = [3, 7, 11]
dilation_sizes = [2, 3, 4]
stride = 1
inp = torch.rand(2, 1, 40)
for kernel_size in kernel_sizes:
for dilation_size in dilation_sizes:
padding = jasper.get_same_padding(kernel_size, stride, dilation_size)
conv = torch.nn.Conv1d(
input_channels, output_channels, kernel_size=kernel_size, dilation=dilation_size, padding=padding
)
out = conv(inp)
assert out.shape == inp.shape
class TestParallelBlock:
@staticmethod
def contrust_jasper_block(**config_kwargs):
config = TestJasperBlock.jasper_base_config(**config_kwargs)
act = jasper.jasper_activations.get(config.pop('activation'))()
block = jasper.JasperBlock(**config, activation=act)
return block
@pytest.mark.unit
def test_blocks_with_same_input_output_channels_sum_residual(self):
blocks = []
in_planes = 8
out_planes = 8
for _ in range(2):
blocks.append(self.contrust_jasper_block(inplanes=in_planes, planes=out_planes))
block = jasper.ParallelBlock(blocks, residual_mode='sum')
x = torch.randn(1, in_planes, 140)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
assert y[0].shape == torch.Size([1, out_planes, 140])
assert ylen[0] == 131
@pytest.mark.unit
def test_blocks_with_different_input_output_channels_sum_residual(self):
blocks = []
in_planes = 8
out_planes = 16
for _ in range(2):
blocks.append(self.contrust_jasper_block(inplanes=in_planes, planes=out_planes))
block = jasper.ParallelBlock(blocks, residual_mode='sum')
x = torch.randn(1, in_planes, 140)
xlen = torch.tensor([131])
with pytest.raises(RuntimeError):
block(([x], xlen))
@pytest.mark.unit
def test_blocks_with_same_input_output_channels_conv_residual(self):
blocks = []
in_planes = 8
out_planes = 8
for _ in range(2):
blocks.append(self.contrust_jasper_block(inplanes=in_planes, planes=out_planes))
block = jasper.ParallelBlock(blocks, residual_mode='conv', in_filters=in_planes, out_filters=out_planes)
x = torch.randn(1, in_planes, 140)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
assert y[0].shape == torch.Size([1, out_planes, 140])
assert ylen[0] == 131
@pytest.mark.unit
def test_blocks_with_different_input_output_channels_conv_residual(self):
blocks = []
in_planes = 8
out_planes = 16
for _ in range(2):
blocks.append(self.contrust_jasper_block(inplanes=in_planes, planes=out_planes))
block = jasper.ParallelBlock(blocks, residual_mode='conv', in_filters=in_planes, out_filters=out_planes)
x = torch.randn(1, in_planes, 140)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
assert y[0].shape == torch.Size([1, out_planes, 140])
assert ylen[0] == 131
@pytest.mark.unit
def test_single_block(self):
in_planes = 8
out_planes = 16
blocks = [self.contrust_jasper_block(inplanes=in_planes, planes=out_planes)]
block = jasper.ParallelBlock(blocks)
x = torch.randn(1, in_planes, 140)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
assert y[0].shape == torch.Size([1, out_planes, 140])
assert ylen[0] == 131
@pytest.mark.unit
def test_tower_dropout(self):
blocks = []
in_planes = 8
out_planes = 8
for _ in range(2):
blocks.append(self.contrust_jasper_block(inplanes=in_planes, planes=out_planes))
block = jasper.ParallelBlock(blocks, aggregation_mode='dropout', block_dropout_prob=1.0)
x = torch.randn(1, in_planes, 140)
xlen = torch.tensor([131])
y, _ = block(([x], xlen))
# Tower dropout is 1.0, meaning that all towers have to be dropped, so only residual remains.
torch.testing.assert_close(y[0], x)