File size: 3,301 Bytes
3139867 | 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 | """Initialize the decoder student by channel-slicing a frozen Vocos decoder.
Training a decoder this small from scratch does not converge to intelligible
audio; slicing the top-left corner of a pretrained Vocos vocoder's weights
does. The Vocos backbone is 512-dim with 4 ConvNeXt blocks and pw 1536; this
takes the first `--dim` channels, the first `--pw` hidden units and the first
3 blocks (dim=40 / pw=120 for the shipped model).
The sliced checkpoint still needs the recovery + z-mix stages (train_decoder.py)
before it is usable -- a fresh slice renders noise, and only the fine-tuning
brings it back.
"""
import torch
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import torch.nn as nn
from microtts.models import DecoderStudent
def slice_linear(src, out_rows, in_cols):
"""Slice a [R, C] weight to [out_rows, in_cols] (top-left corner)."""
return src[:out_rows, :in_cols].contiguous()
def main():
import argparse
ap = argparse.ArgumentParser()
ap.add_argument("--dim", type=int, default=40)
ap.add_argument("--pw", type=int, default=120)
ap.add_argument("--out", default="decoder_init.pt")
args = ap.parse_args()
from vocos import Vocos
v = Vocos.from_pretrained("charactr/vocos-mel-24khz")
sd = v.state_dict()
student = DecoderStudent(dim=args.dim, pw_hidden=args.pw) # 3 blocks
ssd = student.state_dict()
# embed: [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()
# norm
ssd["norm.weight"] = sd["backbone.norm.weight"][:args.dim].contiguous()
ssd["norm.bias"] = sd["backbone.norm.bias"][:args.dim].contiguous()
# blocks 0..2 (student has 3, teacher 4)
for i in range(3):
t = f"backbone.convnext.{i}"
s = f"blocks.{i}"
ssd[f"{s}.gamma"] = sd[f"{t}.gamma"][:args.dim].contiguous()
ssd[f"{s}.dw.weight"] = sd[f"{t}.dwconv.weight"][:args.dim].contiguous()
ssd[f"{s}.dw.bias"] = sd[f"{t}.dwconv.bias"][:args.dim].contiguous()
ssd[f"{s}.norm.weight"] = sd[f"{t}.norm.weight"][:args.dim].contiguous()
ssd[f"{s}.norm.bias"] = sd[f"{t}.norm.bias"][:args.dim].contiguous()
ssd[f"{s}.pw0.weight"] = slice_linear(sd[f"{t}.pwconv1.weight"], args.pw, args.dim)
ssd[f"{s}.pw0.bias"] = sd[f"{t}.pwconv1.bias"][:args.pw].contiguous()
ssd[f"{s}.pw1.weight"] = slice_linear(sd[f"{t}.pwconv2.weight"], args.dim, args.pw)
ssd[f"{s}.pw1.bias"] = sd[f"{t}.pwconv2.bias"][:args.dim].contiguous()
# fnorm
ssd["fnorm.weight"] = sd["backbone.norm.weight"][:args.dim].contiguous()
ssd["fnorm.bias"] = sd["backbone.norm.bias"][:args.dim].contiguous()
# head: [1026, 512] -> [1026, dim]
ssd["head.weight"] = slice_linear(sd["head.out.weight"], 1026, args.dim)
ssd["head.bias"] = sd["head.out.bias"][:1026].contiguous()
student.load_state_dict(ssd)
torch.save(student.state_dict(), args.out)
print(f"saved {args.out}")
n = sum(p.numel() for p in student.parameters())
print(f"student params: {n:,}")
if __name__ == "__main__":
main()
|