File size: 4,749 Bytes
6e802a4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | # test_speaker_init.py: Speaker-table initialisation on a warm start.
import ast
import pathlib
import pytest
ROOT = pathlib.Path(__file__).resolve().parent.parent
SRC = ROOT / "train.py"
@pytest.fixture(scope="module")
def tree():
return ast.parse(SRC.read_text(encoding="utf-8"))
def _flags(tree):
out = set()
for n in ast.walk(tree):
if (isinstance(n, ast.Call)
and getattr(n.func, "attr", "") == "add_argument"
and n.args and isinstance(n.args[0], ast.Constant)):
out.add(n.args[0].value)
return out
def test_both_flags_exist(tree):
f = _flags(tree)
assert "--speaker_init_from" in f
assert "--lr_speaker" in f
def test_speaker_gets_its_own_param_group(tree):
src = SRC.read_text(encoding="utf-8")
assert '"speaker_encoder" in name' in src, "speaker params are not split out"
assert '{"params": spk, "lr": lr_speaker}' in src
def test_warm_start_runs_after_init_from_not_before():
src = SRC.read_text(encoding="utf-8")
load = src.index("load_checkpoint(args.init_from, unwrapped")
copy = src.index("if args.speaker_init_from:")
assert load < copy, "warm-start must come AFTER the init_from load"
def test_warm_start_is_a_plain_copy_under_no_grad():
src = SRC.read_text(encoding="utf-8")
i = src.index("if args.speaker_init_from:")
block = src[i:i + 1400]
assert "torch.no_grad()" in block
assert "W[dst].copy_(W[src])" in block
def test_row_spec_parsing_and_bounds():
def parse(spec, n_rows):
out = []
for pair in spec.split(","):
dst, src = (int(x) for x in pair.strip().split(":"))
if not (0 <= dst < n_rows and 0 <= src < n_rows):
raise ValueError(pair)
out.append((dst, src))
return out
assert parse("4200:1,4201:2", 8192) == [(4200, 1), (4201, 2)]
assert parse(" 4200:1 , 4201:2 ", 8192) == [(4200, 1), (4201, 2)]
with pytest.raises(ValueError):
parse("9000:1", 8192)
with pytest.raises(ValueError):
parse("4200:-1", 8192)
with pytest.raises(ValueError):
parse("4200", 8192)
def test_defaults_preserve_old_behaviour(tree):
for n in ast.walk(tree):
if (isinstance(n, ast.Call)
and getattr(n.func, "attr", "") == "add_argument"
and n.args and isinstance(n.args[0], ast.Constant)
and n.args[0].value in ("--speaker_init_from", "--lr_speaker")):
default = next((k.value for k in n.keywords if k.arg == "default"), None)
assert isinstance(default, ast.Constant)
assert default.value in (None, 0.0), n.args[0].value
def test_lr_speaker_zero_does_not_freeze(tree):
src = SRC.read_text(encoding="utf-8")
assert 'lr_speaker = getattr(args, "lr_speaker", 0.0)' in src, (
"the flag lookup changed — re-check the default is still falsy 0.0")
assert "if lr_speaker and spk:" in src, (
"guard changed — re-check whether --lr_speaker 0 now freezes")
i = src.index("if lr_speaker and spk:")
else_block = src[i:i + 700]
assert "trainable = rest + spk" in else_block, (
"the falsy branch must still be the one that merges spk into the main "
"group; if this changed, this test's premise is stale")
def test_freeze_speaker_flag_exists_and_is_inert_by_default(tree):
f = _flags(tree)
assert "--freeze_speaker" in f
for n in ast.walk(tree):
if (isinstance(n, ast.Call)
and getattr(n.func, "attr", "") == "add_argument"
and n.args and isinstance(n.args[0], ast.Constant)
and n.args[0].value == "--freeze_speaker"):
action = next((k.value.value for k in n.keywords if k.arg == "action"), None)
assert action == "store_true", "must default to off (store_true)"
def test_freeze_runs_before_ddp_wrap_and_before_build_optimizer():
src = SRC.read_text(encoding="utf-8")
freeze = src.index("if getattr(args, \"freeze_speaker\", False):")
ddp = src.index("model = DDP(model, device_ids=[local_rank])")
build = src.index("optimizer = build_optimizer(unwrapped, args)")
assert freeze < ddp, "freeze must precede the DDP wrap"
assert freeze < build, "freeze must precede build_optimizer"
def test_freeze_is_verified_not_assumed():
src = SRC.read_text(encoding="utf-8")
assert "--freeze_speaker did not take" in src
assert "frozen speaker params are in the optimizer" in src
def test_warm_start_still_runs_after_the_freeze():
src = SRC.read_text(encoding="utf-8")
load = src.index("load_checkpoint(args.init_from, unwrapped")
copy = src.index("if args.speaker_init_from:")
assert load < copy
|