| """Convert the official BEVFormer-small checkpoint to a safetensors state_dict for the HF weights mirror. |
| |
| Reads the mmcv checkpoint (as data only: torch.load(weights_only=True)), keeps `state_dict` (drops `optimizer` and the |
| training `meta`), and writes `bevformer_small_epoch_24.safetensors` plus a provenance JSON. Verifies a bit-exact |
| round trip of every tensor. |
| |
| Usage: |
| python -I make_weights_mirror.py <checkpoint.pth> <out_dir> |
| """ |
|
|
| import hashlib |
| import json |
| import sys |
| from pathlib import Path |
|
|
| import torch |
| from safetensors.torch import load_file, save_file |
|
|
| SOURCE_URL = "https://github.com/zhiqi-li/storage/releases/download/v1.0/bevformer_small_epoch_24.pth" |
| EXPECTED_SHA256 = "4ebc5810201ca1e609c29452f07c08bc4d9c9ebfda5d7f8ee1034dd446274bc5" |
|
|
|
|
| def sha256_of(path: Path) -> str: |
| h = hashlib.sha256() |
| with open(path, "rb") as f: |
| for chunk in iter(lambda: f.read(1 << 20), b""): |
| h.update(chunk) |
| return h.hexdigest() |
|
|
|
|
| def main() -> None: |
| ckpt_path = Path(sys.argv[1]) |
| out_dir = Path(sys.argv[2]) |
| out_dir.mkdir(parents=True, exist_ok=True) |
|
|
| src_sha = sha256_of(ckpt_path) |
| if src_sha != EXPECTED_SHA256: |
| raise SystemExit(f"sha256 mismatch: {src_sha} != {EXPECTED_SHA256}") |
|
|
| ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=True) |
| state = ckpt["state_dict"] |
| tensors = {k: v.detach().contiguous().clone() for k, v in state.items()} |
| meta = { |
| "source_url": SOURCE_URL, |
| "source_sha256": src_sha, |
| "source_format": "mmcv checkpoint (state_dict only; optimizer and meta dropped)", |
| "model": "BEVFormer-small (bevformer_small_epoch_24), nuScenes, 24 epochs", |
| } |
| out_file = out_dir / "bevformer_small_epoch_24.safetensors" |
| save_file(tensors, str(out_file), metadata=meta) |
|
|
| back = load_file(str(out_file)) |
| assert back.keys() == tensors.keys() |
| for k, v in tensors.items(): |
| assert back[k].dtype == v.dtype and back[k].shape == v.shape and torch.equal(back[k], v), k |
|
|
| n_params = sum(v.numel() for v in tensors.values() if v.is_floating_point()) |
| dtypes = sorted({str(v.dtype) for v in tensors.values()}) |
| prov = dict(meta) |
| prov.update( |
| { |
| "num_tensors": len(tensors), |
| "num_float_elements": n_params, |
| "dtypes": dtypes, |
| "safetensors_file": out_file.name, |
| "safetensors_sha256": sha256_of(out_file), |
| "safetensors_bytes": out_file.stat().st_size, |
| "epoch": ckpt.get("meta", {}).get("epoch"), |
| "iter": ckpt.get("meta", {}).get("iter"), |
| } |
| ) |
| (out_dir / "provenance.json").write_text(json.dumps(prov, indent=2) + "\n") |
| print(json.dumps(prov, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
