File size: 3,599 Bytes
76a9daf | 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 | """Initialize the decoder student by channel-slicing a frozen pretrained vocoder.
Training a decoder this small from scratch does not converge to intelligible
audio; slicing the top-left corner of a pretrained neural vocoder's decoder
weights does. Pass any ConvNeXt-style vocoder state dict with --vocoder-ckpt;
this takes the first `--dim` channels, the first `--pw` hidden units and the
first `--blocks` blocks (dim 130 / pw 337 / blocks 4 for the shipped model).
The sliced checkpoint still needs decoder training (train_decoder.py) before
it is usable: a fresh slice renders noise, and only fine-tuning brings it back.
python init_decoder.py --dim 130 --pw 337 --blocks 4 \
--vocoder-ckpt <vocos-mel-24khz state dict> --out decoder_init.pt
"""
import sys
from pathlib import Path
import torch
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from microtts.models import DecoderStudent
def main():
import argparse
ap = argparse.ArgumentParser()
ap.add_argument("--dim", type=int, default=130)
ap.add_argument("--pw", type=int, default=337)
ap.add_argument("--blocks", type=int, default=4)
ap.add_argument("--vocoder-ckpt", default=None,
help="state dict / checkpoint of a ConvNeXt-style neural vocoder")
ap.add_argument("--out", default="decoder_init.pt")
args = ap.parse_args()
if args.vocoder_ckpt:
obj = torch.load(args.vocoder_ckpt, map_location="cpu", weights_only=False)
sd = obj.state_dict() if hasattr(obj, "state_dict") else obj
else:
raise ImportError("pass --vocoder-ckpt pointing at a ConvNeXt-style vocoder state dict")
student = DecoderStudent(dim=args.dim, pw_hidden=args.pw, blocks=args.blocks)
ssd = student.state_dict()
# embedding conv: [512, 100, 7] -> [dim, 100, 7]
ssd["embed.weight"] = sd["backbone.embed.weight"][:args.dim].contiguous()
ssd["embed.bias"] = sd["backbone.embed.bias"][:args.dim].contiguous()
ssd["norm.weight"] = sd["backbone.norm.weight"][:args.dim].contiguous()
ssd["norm.bias"] = sd["backbone.norm.bias"][:args.dim].contiguous()
for i in range(args.blocks):
src = f"backbone.convnext.{i}"
dst = f"blocks.{i}"
ssd[f"{dst}.dw.weight"] = sd[f"{src}.dwconv.weight"][:args.dim].contiguous()
ssd[f"{dst}.dw.bias"] = sd[f"{src}.dwconv.bias"][:args.dim].contiguous()
ssd[f"{dst}.norm.weight"] = sd[f"{src}.norm.weight"][:args.dim].contiguous()
ssd[f"{dst}.norm.bias"] = sd[f"{src}.norm.bias"][:args.dim].contiguous()
# pointwise slices: pw0 [pw, dim], pw1 [dim, pw]
ssd[f"{dst}.pw0.weight"] = sd[f"{src}.pwconv1.weight"][:args.pw, :args.dim].contiguous()
ssd[f"{dst}.pw0.bias"] = sd[f"{src}.pwconv1.bias"][:args.pw].contiguous()
ssd[f"{dst}.pw1.weight"] = sd[f"{src}.pwconv2.weight"][:args.dim, :args.pw].contiguous()
ssd[f"{dst}.pw1.bias"] = sd[f"{src}.pwconv2.bias"][:args.dim].contiguous()
ssd[f"{dst}.gamma"] = sd[f"{src}.gamma"][:args.dim].contiguous()
# final norm + head
ssd["fnorm.weight"] = sd["backbone.final_layer_norm.weight"][:args.dim].contiguous()
ssd["fnorm.bias"] = sd["backbone.final_layer_norm.bias"][:args.dim].contiguous()
ssd["head.weight"] = sd["head.out.weight"][:, :args.dim].contiguous()
ssd["head.bias"] = sd["head.out.bias"].contiguous()
student.load_state_dict(ssd)
torch.save(student.state_dict(), args.out)
n = sum(p.numel() for p in student.parameters())
print(f"saved {args.out} | student params: {n:,}")
if __name__ == "__main__":
main()
|