diff --git a/funasr/models/normalize/utterance_mvn.py b/funasr/models/normalize/utterance_mvn.py index 9128c6d044..8c8353e698 100644 --- a/funasr/models/normalize/utterance_mvn.py +++ b/funasr/models/normalize/utterance_mvn.py @@ -81,9 +81,10 @@ def utterance_mvn( x -= mean if norm_vars: + x = x.masked_fill(make_pad_mask(ilens, x, 1), 0.0) var = x.pow(2).sum(dim=1, keepdim=True) / ilens_ std = torch.clamp(var.sqrt(), min=eps) - x = x / std.sqrt() + x = x / std return x, ilens else: if norm_vars: diff --git a/tests/test_utterance_mvn.py b/tests/test_utterance_mvn.py new file mode 100644 index 0000000000..f74201cd03 --- /dev/null +++ b/tests/test_utterance_mvn.py @@ -0,0 +1,268 @@ +"""Run with: python -m pytest tests/test_utterance_mvn.py -q (CPU and CUDA).""" + +from pathlib import Path + +import pytest +import soundfile as sf +import torch +from omegaconf import OmegaConf + +from funasr import AutoModel +from funasr.frontends.wav_frontend import WavFrontend +from funasr.models.normalize.utterance_mvn import UtteranceMVN, utterance_mvn +from funasr.models.transformer.model import Transformer +from funasr.tokenizer.char_tokenizer import CharTokenizer +from funasr.utils.misc import prepare_model_dir + + +@pytest.fixture(params=["cpu", "cuda"]) +def device(request): + if request.param == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA is unavailable") + return request.param + + +def features(device, dtype=torch.float32): + generator = torch.Generator().manual_seed(17) + x = torch.randn(2, 17, 3, generator=generator, dtype=dtype) * 2 + 5 + return x.to(device), torch.tensor([11, 17], device=device) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize( + "norm_means,norm_vars", [(True, True), (True, False), (False, True), (False, False)] +) +def test_formula_and_switches(device, dtype, norm_means, norm_vars): + x, lengths = features(device, dtype) + actual, actual_lengths = utterance_mvn(x.clone(), lengths, norm_means, norm_vars) + assert actual.dtype == dtype + assert actual.device == x.device + assert actual_lengths is lengths + tolerance = { + torch.float64: 1e-12, + torch.float32: 2e-6, + torch.float16: 0.02, + torch.bfloat16: 0.15, + }[dtype] + for i, length in enumerate(lengths.tolist()): + # Independent oracle: unpadded per-utterance torch.std, evaluated in double. + valid = x[i, :length].double() + mean = valid.mean(0) + expected = valid - mean if norm_means else valid + if norm_vars: + expected = expected / valid.std(0, unbiased=False).clamp_min(1e-20) + torch.testing.assert_close( + actual[i, :length].double(), expected, rtol=tolerance, atol=tolerance + ) + if norm_means and norm_vars: + torch.testing.assert_close( + actual[i, :length].double().var(0, unbiased=False), + torch.ones(3, device=device, dtype=torch.float64), + rtol=tolerance, + atol=tolerance, + ) + if norm_means and not norm_vars: + # Preserve historical mean-only padding. + torch.testing.assert_close( + actual[i, length:].double(), + -mean.expand(17 - length, -1), + rtol=tolerance, + atol=tolerance, + ) + else: + assert torch.count_nonzero(actual[i, length:]) == 0 + + +def test_batch_length_and_order_invariance(device): + x, lengths = features(device) + layer = UtteranceMVN(norm_vars=True) + batch, _ = layer(x.clone(), lengths) + single, inferred_lengths = layer(x[:1, :11].clone()) + reverse, _ = layer(x.flip(0), lengths.flip(0)) + torch.testing.assert_close(single[0], batch[0, :11], rtol=1e-6, atol=1e-6) + torch.testing.assert_close(reverse.flip(0), batch, rtol=0, atol=0) + assert inferred_lengths.tolist() == [11] + + +@pytest.mark.parametrize( + "norm_means,norm_vars,expected", + [ + (True, False, [-2.0, 2.0, -3.0, -3.0]), + (False, True, [0.5, 2.5, 0.0, 0.0]), + (False, False, [1.0, 5.0, 0.0, 0.0]), + ], +) +def test_other_branches_keep_existing_values(device, norm_means, norm_vars, expected): + x = torch.tensor([[[1.0], [5.0], [0.0], [0.0]]], device=device) + actual, _ = utterance_mvn(x, torch.tensor([2], device=device), norm_means, norm_vars) + torch.testing.assert_close(actual.flatten(), x.new_tensor(expected), rtol=0, atol=0) + + +@pytest.mark.parametrize("padding", [0, 2]) +def test_mean_variance_values(device, padding): + x = torch.tensor([1.0, 5.0] + [0.0] * padding, device=device).view(1, -1, 1) + actual, _ = UtteranceMVN(norm_vars=True)(x, torch.tensor([2], device=device)) + expected = x.new_tensor([-1.0, 1.0] + [0.0] * padding) + torch.testing.assert_close(actual.flatten(), expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("norm_means", [True, False]) +def test_eps_clamps_standard_deviation(device, norm_means): + x = torch.tensor([[[1.0], [3.0]]], device=device, dtype=torch.float64) + actual, _ = utterance_mvn(x.clone(), norm_means=norm_means, norm_vars=True, eps=4.0) + expected = (x - 2.0 if norm_means else x) / 4.0 + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_corrected_gradients(device): + x, lengths = features(device, torch.float64) + x.requires_grad_() + reference = x.detach().clone().requires_grad_() + actual, _ = UtteranceMVN(norm_vars=True)(x, lengths) + weights = torch.linspace(0.1, 1.3, x.numel(), device=device).reshape_as(x) + (actual * weights).sum().backward() + loss = 0 + for i, length in enumerate(lengths.tolist()): + valid = reference[i, :length] + normalized = (valid - valid.mean(0)) / valid.std(0, unbiased=False) + loss = loss + (normalized * weights[i, :length]).sum() + loss.backward() + assert torch.isfinite(x.grad).all() + assert torch.count_nonzero(x.grad[0, :11]) > 0 + assert torch.count_nonzero(x.grad[0, 11:]) == 0 + torch.testing.assert_close(x.grad, reference.grad, rtol=1e-9, atol=1e-9) + + +def model_config(normalize_conf): + return dict( + model="Transformer", + input_size=80, + vocab_size=9, + encoder="TransformerEncoder", + encoder_conf=dict( + output_size=32, + attention_heads=2, + linear_units=64, + num_blocks=1, + input_layer="linear", + dropout_rate=0.0, + positional_dropout_rate=0.0, + attention_dropout_rate=0.0, + ), + decoder="TransformerDecoder", + decoder_conf=dict( + attention_heads=2, + linear_units=64, + num_blocks=1, + dropout_rate=0.0, + positional_dropout_rate=0.0, + self_attention_dropout_rate=0.0, + src_attention_dropout_rate=0.0, + ), + ctc_weight=0.5, + normalize="UtteranceMVN", + normalize_conf=normalize_conf, + frontend="WavFrontend", + frontend_conf=dict(fs=16000, dither=0.0), + tokenizer="CharTokenizer", + tokenizer_conf=dict( + token_list=["", "", "", "你", "好", "语", "音", "识", ""] + ), + ) + + +def test_saved_config_and_checkpoint_round_trip(tmp_path): + normalize_conf = dict(norm_means=True, norm_vars=True) + config = model_config(normalize_conf) + original = Transformer(**config).eval() + # Use the training config writer and the local AutoModel config/weight loader. + prepare_model_dir(output_dir=str(tmp_path), **config) + torch.save({"state_dict": original.state_dict()}, tmp_path / "model.pt") + saved = OmegaConf.load(tmp_path / "config.yaml") + assert dict(saved.normalize_conf) == normalize_conf + restored = AutoModel( + model=str(tmp_path), device="cpu", ncpu=1, disable_update=True, ignore_init_mismatch=False + ).model + assert restored.normalize.state_dict() == {} + for key, value in original.state_dict().items(): + torch.testing.assert_close(restored.state_dict()[key], value, rtol=0, atol=0) + x = torch.randn(2, 17, 80) + lengths = torch.tensor([11, 17]) + with torch.no_grad(): + expected, _ = original.encode(x.clone(), lengths) + actual, _ = restored.encode(x.clone(), lengths) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("norm_vars", [False, True]) +def test_real_audio_transformer_inference(device, norm_vars): + """Real loader/frontend/encoder/CTC; small random model, no downloads or training.""" + audio, fs = sf.read( + Path(__file__).resolve().parents[1] / "runtime/funasr_api/asr_example.wav", dtype="float32" + ) + short = audio[: 2 * fs].copy() + config = model_config(dict(norm_means=True, norm_vars=norm_vars)) + with torch.random.fork_rng(devices=[]): + torch.manual_seed(2026) + model = Transformer(**config).to(device).eval() + frontend = WavFrontend(**config["frontend_conf"]) + tokenizer = CharTokenizer(**config["tokenizer_conf"]) + captured = {} + + def frontend_hook(module, args, output): + captured["frontend_device"] = output[0].device.type + + def raw_hook(module, args): + captured["raw"] = args[0].detach().clone() + + def norm_hook(module, args, output): + captured["normalized"] = output[0].detach().clone() + captured["lengths"] = output[1].detach().clone() + + def encoder_hook(module, args, output): + captured["encoded"] = output[0].detach().clone() + + handles = [ + frontend.register_forward_hook(frontend_hook), + model.normalize.register_forward_pre_hook(raw_hook), + model.normalize.register_forward_hook(norm_hook), + model.encoder.register_forward_hook(encoder_hook), + ] + outputs = [] + try: + for inputs in [[short], [short, audio], [audio, short]]: + with torch.no_grad(): + results, _ = model.inference( + inputs, + batch_size=2, + frontend=frontend, + tokenizer=tokenizer, + device=device, + fs=fs, + ) + assert len(results) == len(inputs) + assert all("text" in result for result in results) + outputs.append(dict(captured)) + finally: + for handle in handles: + handle.remove() + single, batch, reverse = outputs + n = int(single["lengths"][0]) + assert batch["lengths"].tolist() == [198, 553] + assert next(model.parameters()).device.type == device + assert all(output["frontend_device"] == "cpu" for output in outputs) + for name in ["raw", "normalized", "encoded"]: + assert batch[name].device.type == device + torch.testing.assert_close(single[name][0, :n], batch[name][0, :n], rtol=0, atol=1e-5) + torch.testing.assert_close(batch[name][0, :n], reverse[name][1, :n], rtol=0, atol=1e-5) + if norm_vars: + valid = single["raw"][0, :n].double() + expected = (valid - valid.mean(0)) / valid.std(0, unbiased=False) + torch.testing.assert_close(batch["normalized"][0, :n].double(), expected, rtol=0, atol=2e-6) + torch.testing.assert_close( + batch["normalized"][0, :n].var(0, unbiased=False), + torch.ones(80, device=device), + rtol=0, + atol=1e-6, + ) + assert torch.count_nonzero(batch["normalized"][0, n:]) == 0