| |
| """Warm-start expansion of the enhancer-small DAC decoder (188M -> >=300M). |
| |
| Takes the released `Scicom-intl/CallEnhancer` decoder checkpoint (3072 ch, 188M, |
| 500k steps) and DEEPENS it by appending extra ResidualUnits inside the widest |
| DecoderBlocks, so all 500k steps of training are kept (the neucodec-44k |
| `extend_decoder.py` trick, adapted to a CNN pyramid). |
| |
| Why not "duplicate the last layer" literally: the DAC decoder is not a |
| homogeneous transformer stack — widths halve per block (3072->1536->768->384-> |
| 192->96), so the literal last layers are tiny (a 96-ch ResidualUnit is ~0.07M). |
| The parameter mass lives in the wide early blocks: one ResidualUnit at 1536 ch |
| is ~18.9M (@400 Hz frame rate — cheap compute) and one at 768 ch is ~4.7M |
| (@2 kHz). The default `1536:9,768:5` takes 188M -> ~381M |
| (9x18.9 + 5x4.7 = +193M), i.e. slightly over 2x the released decoder. |
| |
| Why duplicate + ZERO-GATE (not random init, not plain copy): |
| * ResidualUnits compute `x + f(x)`. Zeroing the new unit's output 1x1 conv |
| (weight_norm g + bias) makes it an exact identity, so the expanded model is |
| FUNCTION-IDENTICAL to the 500k checkpoint at step 0 (`--verify` proves it). |
| The 500k-strength discriminator (also carried over) therefore stays balanced |
| instead of hammering a suddenly-degraded generator. |
| * The rest of the new unit is copied from the trained same-dilation unit in |
| its block, so when the gate opens, gradients start from trained filter |
| directions rather than noise. |
| |
| Usage (on the pod, dac importable): |
| hf download Scicom-intl/CallEnhancer decoder_callcentre/decoder_train_state.pt --local-dir ck |
| python expand_decoder.py --ckpt ck/decoder_callcentre/decoder_train_state.pt \ |
| --extra-res 1536:6 --out /Sidon/decoder_callcentre/last.pt --verify |
| |
| # then continue training with the SAME graph: |
| python train_decoder_callcentre.py --dec-channels 3072 --extra-res 1536:6 \ |
| --out /Sidon/decoder_callcentre --steps 800000 ... |
| |
| The output is a trainer-resumable `last.pt`: expanded decoder + the 500k |
| discriminator and its optimizer state; the GENERATOR optimizer is dropped (its |
| param groups changed) — the trainer starts a fresh opt_g, so use a warmup/lower |
| LR for the first ~5-10k steps. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
|
|
| import torch |
|
|
| import dac |
|
|
| RATES = [8, 5, 4, 3, 2] |
| DILATIONS = (1, 3, 9) |
| N_BASE_UNITS = 3 |
| FIRST_UNIT_IDX = 2 |
|
|
|
|
| def parse_extra_res(spec: str) -> dict[int, int]: |
| """'1536:6,768:4' -> {1536: 6, 768: 4} (block output channels -> extra units).""" |
| out: dict[int, int] = {} |
| for part in (spec or "").split(","): |
| part = part.strip() |
| if not part: |
| continue |
| ch, n = part.split(":") |
| out[int(ch)] = int(n) |
| return out |
|
|
|
|
| def build_decoder(input_channel: int, channels: int, extra_res: dict[int, int] | None = None, |
| rates=RATES): |
| """Stock DAC decoder graph + `extra_res[out_ch]` extra ResidualUnits appended at |
| the END of each matching DecoderBlock's Sequential — existing state-dict keys are |
| untouched, new units land at `model.<i>.block.{5,6,...}`. Random init (callers |
| either load a full expanded checkpoint, or warm-start via `expand_state`).""" |
| dec = dac.model.dac.Decoder(input_channel=input_channel, channels=channels, rates=list(rates)) |
| extra_res = extra_res or {} |
| used = set() |
| for i in range(len(rates)): |
| out_ch = channels // 2 ** (i + 1) |
| n_extra = extra_res.get(out_ch, 0) |
| if n_extra <= 0: |
| continue |
| used.add(out_ch) |
| blk = dec.model[1 + i] |
| for k in range(n_extra): |
| unit = dac.model.dac.ResidualUnit(out_ch, dilation=DILATIONS[k % len(DILATIONS)]) |
| blk.block.add_module(str(len(blk.block)), unit) |
| unknown = set(extra_res) - used |
| if unknown: |
| raise ValueError(f"extra_res channels {sorted(unknown)} do not match any " |
| f"DecoderBlock output width of a {channels}-ch decoder") |
| return dec |
|
|
|
|
| def _zero_output_gate(unit) -> None: |
| """Zero the unit's final 1x1 conv (weight-norm g + bias) => unit is identity.""" |
| conv = unit.block[-1] |
| with torch.no_grad(): |
| if hasattr(conv, "weight_g"): |
| conv.weight_g.zero_() |
| elif hasattr(conv, "parametrizations"): |
| conv.parametrizations.weight.original0.zero_() |
| else: |
| conv.weight.zero_() |
| if conv.bias is not None: |
| conv.bias.zero_() |
|
|
|
|
| def expand_state(dec, old_sd: dict, channels: int, extra_res: dict[int, int], |
| rates=RATES) -> None: |
| """Load the trained (unexpanded) `old_sd` into the expanded `dec`, then init each |
| appended unit as a COPY of its block's trained same-dilation unit with the output |
| gate zeroed (exact identity => function-preserving warm start).""" |
| missing, unexpected = dec.load_state_dict(old_sd, strict=False) |
| if unexpected: |
| raise RuntimeError(f"checkpoint keys not in expanded graph: {unexpected[:5]} ...") |
| base_keys = set(old_sd) |
| for key in missing: |
| if ".block." not in key or int(key.split(".block.")[1].split(".")[0]) < FIRST_UNIT_IDX + N_BASE_UNITS: |
| raise RuntimeError(f"missing non-appended key after load: {key}") |
| for i in range(len(rates)): |
| out_ch = channels // 2 ** (i + 1) |
| n_extra = extra_res.get(out_ch, 0) |
| blk = dec.model[1 + i] |
| for k in range(n_extra): |
| src = blk.block[FIRST_UNIT_IDX + (k % N_BASE_UNITS)] |
| new = blk.block[FIRST_UNIT_IDX + N_BASE_UNITS + k] |
| new.load_state_dict(src.state_dict()) |
| _zero_output_gate(new) |
| assert base_keys <= set(dec.state_dict()) |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) |
| ap.add_argument("--ckpt", required=True, |
| help="decoder_train_state.pt / last.pt / decoder_only.pt (needs 'decoder' key)") |
| ap.add_argument("--extra-res", default="1536:9,768:5", |
| help="extra ResidualUnits per block as 'out_ch:n,...' " |
| "[1536:9,768:5 => 188M -> ~381M; 1536:6 => ~301M]") |
| ap.add_argument("--out", required=True, help="output trainer-resumable last.pt") |
| ap.add_argument("--input-channel", type=int, default=1024) |
| ap.add_argument("--verify", action="store_true", |
| help="also build the base decoder and assert identical outputs") |
| ap.add_argument("--set-step", type=int, default=None, |
| help="override the stored step (default: keep, e.g. 500000)") |
| a = ap.parse_args() |
|
|
| ck = torch.load(a.ckpt, map_location="cpu") |
| old_sd = ck["decoder"] if "decoder" in ck else ck |
| channels = ck.get("dec_channels", 3072) |
| if ck.get("extra_res"): |
| raise SystemExit(f"{a.ckpt} is already expanded (extra_res={ck['extra_res']})") |
| extra = parse_extra_res(a.extra_res) |
|
|
| dec = build_decoder(a.input_channel, channels, extra) |
| expand_state(dec, old_sd, channels, extra) |
| n_old = sum(v.numel() for v in old_sd.values()) |
| n_new = sum(p.numel() for p in dec.parameters()) |
| print(f"[expand] {channels}ch +{a.extra_res}: {n_old/1e6:.1f}M -> {n_new/1e6:.1f}M params") |
|
|
| if a.verify: |
| base = build_decoder(a.input_channel, channels, None) |
| base.load_state_dict(old_sd) |
| base.eval(); dec.eval() |
| with torch.no_grad(): |
| x = torch.randn(1, a.input_channel, 50) |
| d = (base(x) - dec(x)).abs().max().item() |
| print(f"[verify] max |base - expanded| on random input: {d:.3e}") |
| assert d < 1e-5, "expanded decoder is NOT function-preserving" |
|
|
| out = {"step": a.set_step if a.set_step is not None else ck.get("step", 0), |
| "decoder": dec.state_dict(), |
| "dec_channels": channels, |
| "extra_res": a.extra_res} |
| for key in ("disc", "opt_d"): |
| if key in ck: |
| out[key] = ck[key] |
| torch.save(out, a.out) |
| print(f"[save] {a.out} (step={out['step']}, disc={'disc' in out}, opt_d={'opt_d' in out}; " |
| f"opt_g intentionally dropped -> trainer starts a fresh generator optimizer)") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
