File size: 3,236 Bytes
ad7a3dd | 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 | """Offline graph/weight inspection; TensorFlow is confined to this tooling."""
import argparse
import hashlib
import json
import os
from pathlib import Path
os.environ.setdefault("TF_ENABLE_ONEDNN_OPTS", "0")
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")
import numpy as np
import tensorflow as tf
from tensorflow.core.protobuf import saved_model_pb2
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--source", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
args = parser.parse_args()
args.output.mkdir(parents=True, exist_ok=False)
identity = json.loads((args.source / ".source.json").read_text())
for name, expected in identity["files"].items():
if hashlib.sha256((args.source / name).read_bytes()).hexdigest() != expected:
raise ValueError("Source Google artifact changed: " + name)
saved = saved_model_pb2.SavedModel()
saved.ParseFromString((args.source / "saved_model.pb").read_bytes())
graph = next(g for g in saved.meta_graphs if not g.meta_info_def.tags)
(args.output / "graph.pb").write_bytes(graph.graph_def.SerializeToString())
ops = []
constants = {}
for node in graph.graph_def.node:
attrs = {}
for name, value in node.attr.items():
field = value.WhichOneof("value")
if field in ("s", "i", "f", "b"):
v = getattr(value, field)
attrs[name] = v.decode(errors="replace") if isinstance(v, bytes) else v
elif field == "list":
attrs[name] = {k: list(getattr(value.list, k)) for k in ("i", "f", "b") if len(getattr(value.list, k))}
elif field == "shape":
attrs[name] = [d.size for d in value.shape.dim]
ops.append({"name": node.name, "op": node.op, "inputs": list(node.input), "attrs": attrs})
if node.op == "Const":
array = tf.make_ndarray(node.attr["value"].tensor)
if array.dtype.kind in "fiu" and array.size:
constants[node.name] = array
reader = tf.train.load_checkpoint(str(args.source / "variables/variables"))
variables = {name: reader.get_tensor(name) for name in sorted(reader.get_variable_to_shape_map())}
np.savez(args.output / "variables.npz", **variables)
np.savez(args.output / "constants.npz", **constants)
(args.output / "ops.json").write_text(json.dumps(ops, indent=2) + "\n")
report = {"tensorflow": tf.__version__, "numpy": np.__version__, "source": identity,
"tags": [list(m.meta_info_def.tags) for m in saved.meta_graphs],
"variables": {k: list(v.shape) for k, v in variables.items()},
"signatures": {name: {"inputs": {k: v.name for k, v in s.inputs.items()},
"outputs": {k: v.name for k, v in s.outputs.items()}} for name, s in graph.signature_def.items()},
"op_counts": {op: sum(n["op"] == op for n in ops) for op in sorted({n["op"] for n in ops})}}
(args.output / "report.json").write_text(json.dumps(report, indent=2) + "\n")
print(json.dumps(report, indent=2))
if __name__ == "__main__":
main()
|