File size: 4,477 Bytes
f49837b | 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 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 | #!/usr/bin/env python3
"""Create an inference-only MMEngine checkpoint from a trusted training checkpoint.
The source checkpoints used for this release contain MMEngine runtime objects,
the full training configuration, optimizer/runtime state, and local filesystem
paths. This script loads those trusted files with an explicit safe-global
allowlist and writes only:
* ``state_dict`` tensors on CPU
* minimal, non-sensitive dataset metadata
* the source epoch and a format version
Never use this script on an untrusted checkpoint.
"""
from __future__ import annotations
import argparse
import builtins
import hashlib
import zipfile
from collections import OrderedDict
from pathlib import Path
import numpy as np
import torch
from mmengine.logging.history_buffer import HistoryBuffer
_DTYPES = (
np.float16,
np.float32,
np.float64,
np.int8,
np.int16,
np.int32,
np.int64,
np.uint8,
np.uint16,
np.uint32,
np.uint64,
np.bool_,
)
_SAFE_GLOBALS = [
HistoryBuffer,
(np._core.multiarray._reconstruct, "numpy.core.multiarray._reconstruct"),
(np._core.multiarray.scalar, "numpy.core.multiarray.scalar"),
np.ndarray,
np.dtype,
*{type(np.dtype(dtype)) for dtype in _DTYPES},
builtins.getattr,
]
_FORBIDDEN_MARKERS = (
b"/data",
b"HistoryBuffer",
b"message_hub",
b"optimizer",
b"experiment_name",
)
def _sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as handle:
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def _pickle_payload(path: Path) -> bytes:
with zipfile.ZipFile(path) as archive:
pickle_names = [name for name in archive.namelist() if name.endswith("/data.pkl")]
if len(pickle_names) != 1:
raise RuntimeError(f"Expected one data.pkl member, found {pickle_names}")
return archive.read(pickle_names[0])
def sanitize(source: Path, destination: Path) -> None:
torch.serialization.add_safe_globals(_SAFE_GLOBALS)
checkpoint = torch.load(source, map_location="cpu", weights_only=True)
if not isinstance(checkpoint, dict) or "state_dict" not in checkpoint:
raise TypeError("Source is not an MMEngine checkpoint with a state_dict")
source_state = checkpoint["state_dict"]
if not isinstance(source_state, dict):
raise TypeError("state_dict must be a mapping")
if not source_state or not all(torch.is_tensor(value) for value in source_state.values()):
raise TypeError("state_dict must contain tensors only")
state_dict = OrderedDict(
(name, tensor.detach().cpu().contiguous())
for name, tensor in source_state.items()
)
source_meta = checkpoint.get("meta") or {}
source_epoch = int(source_meta.get("epoch", 0))
clean_checkpoint = {
"state_dict": state_dict,
"meta": {
"dataset_meta": {
"classes": ("nodule",),
"palette": [(220, 20, 60)],
},
"source_epoch": source_epoch,
"format_version": 1,
},
}
destination.parent.mkdir(parents=True, exist_ok=True)
torch.save(clean_checkpoint, destination)
verified = torch.load(destination, map_location="cpu", weights_only=True)
verified_state = verified.get("state_dict", {})
if list(verified_state) != list(state_dict):
raise RuntimeError("State-dict keys changed during serialization")
for name, tensor in state_dict.items():
if not torch.equal(tensor, verified_state[name]):
raise RuntimeError(f"Tensor changed during serialization: {name}")
payload = _pickle_payload(destination)
leaked = [marker.decode("ascii") for marker in _FORBIDDEN_MARKERS if marker in payload]
if leaked:
raise RuntimeError(f"Sanitized checkpoint still contains forbidden markers: {leaked}")
print(f"source_sha256={_sha256(source)}")
print(f"output_sha256={_sha256(destination)}")
print(f"source_epoch={source_epoch}")
print(f"tensor_count={len(state_dict)}")
print(f"parameter_numel={sum(tensor.numel() for tensor in state_dict.values())}")
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("source", type=Path)
parser.add_argument("destination", type=Path)
args = parser.parse_args()
sanitize(args.source, args.destination)
if __name__ == "__main__":
main()
|