File size: 8,672 Bytes
0993d8a | 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 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 | #!/usr/bin/env python3
"""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) # DAC's per-block dilation stack; appended units cycle it
N_BASE_UNITS = 3 # ResidualUnits per DecoderBlock in stock DAC
FIRST_UNIT_IDX = 2 # DecoderBlock.block = [Snake1d, WNConvTranspose1d, RU, RU, RU]
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] # model[0] is the input conv
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"): # torch old-style weight_norm
conv.weight_g.zero_()
elif hasattr(conv, "parametrizations"): # new-style parametrization
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)] # dilation-matched trained unit
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"): # keep the 500k discriminator (+ its optimizer);
if key in ck: # opt_g is dropped — decoder param groups changed
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()
|