| |
| """Writes src/weights.rs: the vendored prototype_model.pt as Q24 integers, read WITHOUT torch. |
| |
| The checkpoint is a torch zip; data.pkl is unpickled with a restricted Unpickler that resolves only |
| collections.OrderedDict, torch._utils._rebuild_tensor_v2 and torch.FloatStorage, and each tensor's bytes are |
| read from <archive>/data/<key> as little-endian float32. Each weight becomes floor(w * 2^24 + 1/2), exactly. |
| The state dict's keys and its _metadata module tree must be those of PrototypeGNN(6, 64, 2, 3) built with the |
| fallback GCNConv of prototype_gnn.py:17-24 (conv*.lin, no aggr_module); the Rust names follow those keys. |
| --check regenerates in memory and diffs. |
| """ |
| import collections |
| import hashlib |
| import io |
| import math |
| import pathlib |
| import pickle |
| import struct |
| import sys |
| import zipfile |
| from fractions import Fraction |
|
|
| ROOT = pathlib.Path(__file__).resolve().parent.parent |
| CHECKPOINT = ROOT / "vendor" / "fnns" / "models" / "prototype_model.pt" |
| CHECKPOINT_SHA256 = "b4df990138ae215bbdda7c1fcf4902bc0a8cb72d9b7d8944eb965d0fc737dd22" |
| WEIGHTS_RS = ROOT / "src" / "weights.rs" |
| FRAC_BITS = 24 |
| VALUES_PER_LINE = 16 |
|
|
| |
| TENSORS = [ |
| ("prototypes", "PROTOTYPES", (2, 3, 64)), |
| ("conv1.lin.weight", "CONV1_W", (64, 6)), |
| ("conv1.lin.bias", "CONV1_B", (64,)), |
| ("conv2.lin.weight", "CONV2_W", (64, 64)), |
| ("conv2.lin.bias", "CONV2_B", (64,)), |
| ("edge_projection.0.weight", "PROJ0_W", (64, 128)), |
| ("edge_projection.0.bias", "PROJ0_B", (64,)), |
| ("edge_projection.2.weight", "PROJ2_W", (64, 64)), |
| ("edge_projection.2.bias", "PROJ2_B", (64,)), |
| ] |
|
|
| |
| MODULE_TREE = [ |
| "", |
| "conv1", |
| "conv1.lin", |
| "conv2", |
| "conv2.lin", |
| "edge_projection", |
| "edge_projection.0", |
| "edge_projection.1", |
| "edge_projection.2", |
| ] |
|
|
|
|
| class FloatStorageClass: |
| """Stands in for the torch.FloatStorage global; never instantiated.""" |
|
|
|
|
| class Storage: |
| def __init__(self, values): |
| self.values = values |
|
|
|
|
| class Tensor: |
| def __init__(self, shape, values): |
| self.shape = shape |
| self.values = values |
|
|
|
|
| def row_major_strides(shape): |
| strides = [] |
| step = 1 |
| for dim in reversed(shape): |
| strides.append(step) |
| step *= dim |
| return tuple(reversed(strides)) |
|
|
|
|
| def rebuild_tensor_v2(storage, storage_offset, size, stride, requires_grad, backward_hooks, metadata=None): |
| |
| if not isinstance(storage, Storage): |
| raise pickle.UnpicklingError("tensor storage is not a float32 storage") |
| size, stride = tuple(size), tuple(stride) |
| if not all(isinstance(d, int) and d > 0 for d in size): |
| raise pickle.UnpicklingError(f"tensor size {size} is not positive integers") |
| if stride != row_major_strides(size): |
| raise pickle.UnpicklingError(f"tensor stride {stride} is not row-major for size {size}") |
| if storage_offset != 0 or math.prod(size) != len(storage.values): |
| raise pickle.UnpicklingError("tensor does not cover its whole storage") |
| if requires_grad is not False or backward_hooks or metadata is not None: |
| raise pickle.UnpicklingError("tensor carries grad, hooks or metadata") |
| return Tensor(size, storage.values) |
|
|
|
|
| class CheckpointUnpickler(pickle.Unpickler): |
| ALLOWED = { |
| ("collections", "OrderedDict"): collections.OrderedDict, |
| ("torch._utils", "_rebuild_tensor_v2"): rebuild_tensor_v2, |
| ("torch", "FloatStorage"): FloatStorageClass, |
| } |
|
|
| def __init__(self, data, archive, prefix): |
| super().__init__(io.BytesIO(data)) |
| self.archive = archive |
| self.prefix = prefix |
|
|
| def find_class(self, module, name): |
| try: |
| return self.ALLOWED[(module, name)] |
| except KeyError: |
| raise pickle.UnpicklingError(f"global {module}.{name} is not allowed") from None |
|
|
| def persistent_load(self, pid): |
| if not (isinstance(pid, tuple) and len(pid) == 5 and pid[0] == "storage"): |
| raise pickle.UnpicklingError(f"persistent id {pid!r} is not a storage") |
| _, storage_type, key, location, numel = pid |
| if storage_type is not FloatStorageClass or location != "cpu": |
| raise pickle.UnpicklingError(f"storage {key!r} is not a cpu FloatStorage") |
| if not (isinstance(key, str) and key.isdigit() and isinstance(numel, int) and numel > 0): |
| raise pickle.UnpicklingError(f"storage id {pid!r} is malformed") |
| raw = self.archive.read(f"{self.prefix}/data/{key}") |
| if len(raw) != 4 * numel: |
| raise pickle.UnpicklingError(f"storage {key} holds {len(raw)} bytes, expected {4 * numel}") |
| return Storage(struct.unpack(f"<{numel}f", raw)) |
|
|
|
|
| def load_state_dict(path=CHECKPOINT): |
| blob = path.read_bytes() |
| digest = hashlib.sha256(blob).hexdigest() |
| if digest != CHECKPOINT_SHA256: |
| raise SystemExit(f"QUANTIZE|refused=sha256|got={digest}|want={CHECKPOINT_SHA256}") |
| with zipfile.ZipFile(io.BytesIO(blob)) as archive: |
| pickles = [n for n in archive.namelist() if n.endswith("/data.pkl")] |
| if len(pickles) != 1: |
| raise SystemExit("QUANTIZE|refused=archive|reason=expected one data.pkl") |
| prefix = pickles[0][: -len("/data.pkl")] |
| if archive.read(f"{prefix}/byteorder") != b"little": |
| raise SystemExit("QUANTIZE|refused=byteorder") |
| state = CheckpointUnpickler(archive.read(pickles[0]), archive, prefix).load() |
| if not isinstance(state, collections.OrderedDict): |
| raise SystemExit("QUANTIZE|refused=root|reason=not an OrderedDict") |
| if [key for key, _, _ in TENSORS] != list(state.keys()): |
| raise SystemExit(f"QUANTIZE|refused=keys|got={','.join(state.keys())}") |
| tree = list(getattr(state, "_metadata", {}).keys()) |
| if tree != MODULE_TREE: |
| raise SystemExit(f"QUANTIZE|refused=module_tree|got={','.join(tree)}") |
| for key, _, shape in TENSORS: |
| if state[key].shape != shape: |
| raise SystemExit(f"QUANTIZE|refused=shape|key={key}|got={state[key].shape}") |
| return state |
|
|
|
|
| def q24(value): |
| return math.floor(Fraction(value) * (1 << FRAC_BITS) + Fraction(1, 2)) |
|
|
|
|
| def q24_tensors(state=None): |
| """{Rust name: (shape, row-major Q24 values)} of the pinned checkpoint.""" |
| state = load_state_dict() if state is None else state |
| out = {} |
| for key, name, shape in TENSORS: |
| values = [q24(v) for v in state[key].values] |
| if not all(-(1 << 31) <= v < (1 << 31) for v in values): |
| raise SystemExit(f"QUANTIZE|refused=range|key={key}") |
| out[name] = (shape, values) |
| return out |
|
|
|
|
| def render_row(values, indent): |
| lines = [] |
| for start in range(0, len(values), VALUES_PER_LINE): |
| lines.append(indent + ", ".join(str(v) for v in values[start:start + VALUES_PER_LINE]) + ",") |
| return lines |
|
|
|
|
| def render_matrix(values, rows, cols, indent): |
| out = [] |
| for r in range(rows): |
| row = values[r * cols:(r + 1) * cols] |
| if cols <= VALUES_PER_LINE: |
| out.append(indent + "[" + ", ".join(str(v) for v in row) + "],") |
| else: |
| out.append(indent + "[") |
| out.extend(render_row(row, indent + " ")) |
| out.append(indent + "],") |
| return out |
|
|
|
|
| def render(tensors): |
| out = [ |
| f"//! Q24 weights, floor(w * 2^24 + 1/2), of vendor/fnns/models/prototype_model.pt (sha256 {CHECKPOINT_SHA256});", |
| "//! generated by tools/quantize.py, do not edit. Keys and module tree are PrototypeGNN(6, 64, 2, 3) with the", |
| "//! fallback GCNConv of prototype_gnn.py:17-24 (`conv*.lin.weight`, `conv*.lin.bias`).", |
| ] |
| for key, name, shape in TENSORS: |
| _, values = tensors[name] |
| out.append("") |
| if len(shape) == 1: |
| out.append(f"/// `{key}` ({shape[0]}), Q24.") |
| out.append("#[rustfmt::skip]") |
| out.append(f"pub const {name}: [i32; {shape[0]}] = [") |
| out.extend(render_row(values, " ")) |
| out.append("];") |
| elif len(shape) == 2: |
| rows, cols = shape |
| out.append(f"/// `{key}` ({rows} x {cols}, row = output), Q24.") |
| out.append("#[rustfmt::skip]") |
| out.append(f"pub const {name}: [[i32; {cols}]; {rows}] = [") |
| out.extend(render_matrix(values, rows, cols, " ")) |
| out.append("];") |
| else: |
| classes, per_class, dim = shape |
| out.append(f"/// `{key}` ({classes} classes x {per_class} prototypes x {dim}), Q24.") |
| out.append("#[rustfmt::skip]") |
| out.append(f"pub const {name}: [[[i32; {dim}]; {per_class}]; {classes}] = [") |
| for c in range(classes): |
| out.append(" [") |
| block = values[c * per_class * dim:(c + 1) * per_class * dim] |
| out.extend(render_matrix(block, per_class, dim, " ")) |
| out.append(" ],") |
| out.append("];") |
| return "\n".join(out) + "\n" |
|
|
|
|
| def main(): |
| args = sys.argv[1:] |
| if args not in ([], ["--check"]): |
| print("usage: quantize.py [--check]", file=sys.stderr) |
| return 2 |
| text = render(q24_tensors()) |
| values = sum(math.prod(shape) for _, _, shape in TENSORS) |
| if args == ["--check"]: |
| if WEIGHTS_RS.read_text(encoding="utf-8") != text: |
| print(f"QUANTIZE|check=FAIL|file={WEIGHTS_RS.name}") |
| return 1 |
| print(f"QUANTIZE|check=PASS|values={values}|sha256={CHECKPOINT_SHA256}") |
| return 0 |
| WEIGHTS_RS.write_text(text, encoding="utf-8") |
| print(f"QUANTIZE|written={WEIGHTS_RS.name}|values={values}|sha256={CHECKPOINT_SHA256}") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
