File size: 3,138 Bytes
4f5f309 | 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 | """Export the 15 per-level depth graphs (fp32) to work_dir()/depth_build/fp32/level_XX.onnx.
Each level's weights are read from the checkpoint and loaded into a fresh DepthLevel; one level at a time.
.venv/bin/python tools/memguard.py --rss-limit-mb 4000 --min-free-mb 4000 -- \
.venv-tts/bin/python -m tools.maya_export.export_depth
"""
from __future__ import annotations
import torch
from tts.model.depth_v2 import GROUP_OF_LEVEL
from tts.tokens import N_Q
from tools.maya_export.common import ckpt_paths, work_dir
from tools.maya_export.depth_frame import DepthLevel, Q_SERVE
from tools.maya_export.export_backbone import D_MODEL, _export, _quantize_wo8, _save_external
from tools.maya_export.streamed_backbone import _Weights, shard_files
D_DEPTH = 1024
N_LEVELS = 15 # serving q_serve = 16: levels 1..15 (the checkpoint has up to 31)
def ckpt_name(name: str, level: int) -> str:
"""DepthLevel parameter name -> checkpoint key (without the model. prefix)."""
if name.startswith("trunk."):
_, li, rest = name.split(".", 2)
return f"depth.groups.{GROUP_OF_LEVEL[level]}.{li}.{rest}"
if name.startswith("head."):
return f"depth.heads.{level - 1}.weight"
return f"depth.{name}"
def _head_nodes(model) -> list[str]:
"""Names of the node(s) producing the `logits` output (the level head's MatMul): kept in fp32 for wo8_fph."""
names = []
for node in model.graph.node:
if "logits" in node.output:
names.append(node.name)
return names
def export_level(level: int, weights: _Weights, build, precision: str) -> None:
mod = DepthLevel(level).eval()
state = {name: weights.get(ckpt_name(name, level)).float() for name, _ in mod.named_parameters()}
mod.load_state_dict(state, strict=True)
del state
h = torch.randn(1, D_MODEL) * 0.02
codes = torch.zeros(1, Q_SERVE, dtype=torch.long)
mem = torch.zeros(1, 1, D_DEPTH)
tmp = build / f"level_{level:02d}_raw.onnx"
_export(mod, (h, codes, mem), tmp, input_names=["h", "codes", "mem"], output_names=["logits"],
dynamic_axes={"h": {0: "B"}, "codes": {0: "B"}, "mem": {0: "B"}, "logits": {0: "B"}})
import onnx
model = onnx.load(str(tmp))
if precision == "wo8":
model = _quantize_wo8(model)
elif precision == "wo8_fph": # int8 trunks, fp32 head
model = _quantize_wo8(model, exclude=_head_nodes(model))
_save_external(model, build / f"level_{level:02d}.onnx")
tmp.unlink()
tmp.with_suffix(".onnx.data").unlink(missing_ok=True)
del mod
def main() -> None:
import argparse
ap = argparse.ArgumentParser()
ap.add_argument("--precision", choices=("fp32", "wo8", "wo8_fph"), default="fp32")
precision = ap.parse_args().precision
build = work_dir() / "depth_build" / precision
build.mkdir(parents=True, exist_ok=True)
paths = ckpt_paths()
weights = _Weights(shard_files(paths["ckpt"]))
for level in range(1, N_LEVELS + 1):
export_level(level, weights, build, precision)
print(f"level {level:02d} fp32 ok", flush=True)
if __name__ == "__main__":
main()
|