Spaces:
Sleeping
Sleeping
File size: 2,669 Bytes
add24bd | 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 | #!/usr/bin/env python3
"""One-time: extract prebackbone weights from a full YOLO best.pt checkpoint."""
from __future__ import annotations
import argparse
import sys
from pathlib import Path
import torch
# Allow running before vendor is installed
_HERE = Path(__file__).resolve().parent
for base in (_HERE / "vendor", _HERE.parent / "ultralytics"):
if (base / "ultralytics" / "__init__.py").exists():
sys.path.insert(0, str(base))
break
def _clean_state_dict(state: dict) -> dict:
"""Drop thop profiling keys (total_ops, total_params) not in nn.Module."""
skip = ("total_ops", "total_params")
return {k: v for k, v in state.items() if not any(k == s or k.endswith(f".{s}") for s in skip)}
def extract(
full_ckpt: Path,
out_path: Path,
prebackbone_name: str | None = None,
) -> Path:
from ultralytics.nn.tasks import torch_safe_load
ckpt, _ = torch_safe_load(str(full_ckpt))
train_args = ckpt.get("train_args") or {}
name = (prebackbone_name or train_args.get("prebackbone") or "A11_CA").upper()
channels = int(train_args.get("channels", 3) or 3)
model = (ckpt.get("ema") or ckpt.get("model")).float()
pb = getattr(model, "prebackbone", None)
if pb is None:
raise RuntimeError(f"No prebackbone in {full_ckpt}")
state = _clean_state_dict(pb.state_dict())
# Save only tensors + metadata (no YOLO backbone/head)
payload = {
"prebackbone": name,
"channels": channels,
"state_dict": state,
"source_checkpoint": str(full_ckpt.resolve()),
}
out_path.parent.mkdir(parents=True, exist_ok=True)
torch.save(payload, out_path)
n_params = sum(t.numel() for t in payload["state_dict"].values())
size_mb = out_path.stat().st_size / (1024 * 1024)
print(f"Saved {out_path} ({size_mb:.2f} MB, {n_params:,} parameters, type={name})")
return out_path
def main() -> None:
p = argparse.ArgumentParser(description="Extract prebackbone-only weights from best.pt")
p.add_argument(
"--ckpt",
type=Path,
default=_HERE / "weights" / "best.pt",
help="Full YOLO checkpoint (best.pt)",
)
p.add_argument(
"--out",
type=Path,
default=_HERE / "weights" / "prebackbone_a11_ca.pt",
help="Output path for prebackbone-only weights",
)
p.add_argument("--name", type=str, default=None, help="Override prebackbone type (default: from train_args)")
args = p.parse_args()
if not args.ckpt.exists():
raise SystemExit(f"Checkpoint not found: {args.ckpt}")
extract(args.ckpt, args.out, args.name)
if __name__ == "__main__":
main()
|