mirror of
https://github.com/sunnypilot/sunnypilot.git
synced 2026-07-19 23:22:06 +08:00
Armand du Parc Locmaria
551e2f77bf
* modeld: standalone compile script * cleanup * frame skip * rm last op import * dm warp * no graph break * +x compile_dm_warp.py * don't import tg before setting device * compile_modeld exports metadata * update help * namedtuple * lint * Revert "compile_modeld exports metadata" This reverts commit 93c3c223567b4d4a074c9071d7f734c56f5aedcc. * import
60 lines
1.9 KiB
Python
Executable File
60 lines
1.9 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
import sys
|
|
import pathlib
|
|
import codecs
|
|
import pickle
|
|
from typing import Any
|
|
|
|
from tinygrad.nn.onnx import OnnxPBParser
|
|
|
|
def metadata_path_for(onnx_path) -> pathlib.Path:
|
|
p = pathlib.Path(onnx_path)
|
|
return p.parent / (p.stem + '_metadata.pkl')
|
|
|
|
|
|
class MetadataOnnxPBParser(OnnxPBParser):
|
|
def _parse_ModelProto(self) -> dict:
|
|
obj: dict[str, Any] = {"graph": {"input": [], "output": []}, "metadata_props": []}
|
|
for fid, wire_type in self._parse_message(self.reader.len):
|
|
match fid:
|
|
case 7:
|
|
obj["graph"] = self._parse_GraphProto()
|
|
case 14:
|
|
obj["metadata_props"].append(self._parse_StringStringEntryProto())
|
|
case _:
|
|
self.reader.skip_field(wire_type)
|
|
return obj
|
|
|
|
|
|
def get_name_and_shape(value_info: dict[str, Any]) -> tuple[str, tuple[int, ...]]:
|
|
shape = tuple(int(dim) if isinstance(dim, int) else 0 for dim in value_info["parsed_type"].shape)
|
|
name = value_info["name"]
|
|
return name, shape
|
|
|
|
|
|
def get_metadata_value_by_name(model: dict[str, Any], name: str) -> str | Any:
|
|
for prop in model["metadata_props"]:
|
|
if prop["key"] == name:
|
|
return prop["value"]
|
|
return None
|
|
|
|
|
|
if __name__ == "__main__":
|
|
model_path = pathlib.Path(sys.argv[1])
|
|
model = MetadataOnnxPBParser(model_path).parse()
|
|
output_slices = get_metadata_value_by_name(model, 'output_slices')
|
|
assert output_slices is not None, 'output_slices not found in metadata'
|
|
|
|
metadata = {
|
|
'model_checkpoint': get_metadata_value_by_name(model, 'model_checkpoint'),
|
|
'output_slices': pickle.loads(codecs.decode(output_slices.encode(), "base64")),
|
|
'input_shapes': dict(get_name_and_shape(x) for x in model["graph"]["input"]),
|
|
'output_shapes': dict(get_name_and_shape(x) for x in model["graph"]["output"]),
|
|
}
|
|
|
|
metadata_path = metadata_path_for(model_path)
|
|
with open(metadata_path, 'wb') as f:
|
|
pickle.dump(metadata, f)
|
|
|
|
print(f'saved metadata to {metadata_path}')
|