File size: 7,232 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 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 | """One-time TensorFlow reference/constant extraction for the PyTorch port."""
import argparse
import hashlib
import json
import os
import sqlite3
from pathlib import Path
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
os.environ.setdefault("TF_ENABLE_ONEDNN_OPTS", "0")
os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")
import numpy as np
import tensorflow as tf
from tensorflow.core.framework import graph_pb2
from tensorflow.core.protobuf import saved_model_pb2
def constant_value(graph_def, name):
"""Evaluate the original constant-only subgraph, without map control edges."""
nodes = {n.name: n for n in graph_def.node}
selected = set()
def add(key):
key = key.split(":")[0]
if key in selected:
return
node = nodes[key]
if node.op in ("Placeholder", "ReadVariableOp", "VarHandleOp"):
raise ValueError("Requested frontend constant depends on runtime state: " + key)
selected.add(key)
for parent in node.input:
if not parent.startswith("^"):
add(parent)
add(name)
subset = graph_pb2.GraphDef()
for node in graph_def.node:
if node.name in selected:
copied = subset.node.add()
copied.CopyFrom(node)
inputs = [v for v in copied.input if not v.startswith("^")]
del copied.input[:]
copied.input.extend(inputs)
with tf.Graph().as_default() as graph:
tf.import_graph_def(subset, name="")
with tf.compat.v1.Session(graph=graph) as session:
return session.run(graph.get_tensor_by_name(name + ":0"))
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--project", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
args = parser.parse_args()
root = args.project.resolve()
output = args.output.resolve()
output.mkdir(parents=True, exist_ok=False)
source = root / "data/artifacts/google-speech-embedding"
identity = json.loads((source / ".source.json").read_text())
for name, expected in identity["files"].items():
if hashlib.sha256((source / name).read_bytes()).hexdigest() != expected:
raise ValueError("Original artifact changed")
saved = saved_model_pb2.SavedModel()
saved.ParseFromString((source / "saved_model.pb").read_bytes())
meta = next(m for m in saved.meta_graphs if not m.meta_info_def.tags)
window = constant_value(meta.graph_def, "map/while/stft/hann_window/sub_2")
mel = constant_value(meta.graph_def, "map/while/linear_to_mel_weight_matrix")
reader = tf.train.load_checkpoint(str(source / "variables/variables"))
weights = {k: reader.get_tensor(k) for k in sorted(reader.get_variable_to_shape_map())}
key = "embeddings/state/MudpuppyLite/ZeroMeanConv2d/weights"
initial = tf.convert_to_tensor(weights[key])
centered = (initial - tf.reduce_mean(initial, axis=[0, 1], keepdims=True)).numpy()
np.savez(output / "frontend.npz", window=window, mel_matrix=mel, initial_centered_hwio=centered)
np.savez(output / "variables.npz", **weights)
manifest = json.loads((root / "data/derived/development-clean460-v1/manifest.json").read_text())
crops = sorted((r for r in manifest["query_crops"].values() if r["split"] == "dev-select"), key=lambda r: (r["num_samples"], r["crop_id"]))
positions = np.unique(np.linspace(0, len(crops) - 1, 32, dtype=int))
db = sqlite3.connect((root / "data/cache/keyword-clean460-google-v1/features.sqlite3").as_uri() + "?mode=ro&immutable=1", uri=True)
examples = []
for i in positions:
crop = crops[i]
row = db.execute("SELECT pcm16,data,valid_length,sha256 FROM features WHERE crop_id=?", (crop["crop_id"],)).fetchone()
if hashlib.sha256(row[0]).hexdigest() != crop["audio_sha256"] or hashlib.sha256(row[1]).hexdigest() != row[3]:
raise ValueError("Retained feature/crop identity changed")
wave = np.frombuffer(row[0], dtype="<i2").astype(np.float32) / 32768
cached = np.frombuffer(row[1], dtype="<f4").copy().reshape(row[2], 96)
examples.append(("speech-" + crop["crop_id"], wave, cached))
db.close()
rng = np.random.default_rng(20261011)
for length in (8000, 8001, 8960, 16000, 31999, 32000):
examples.append((f"silence-{length}", np.zeros(length, dtype=np.float32), None))
for length in (8000, 16000, 32000):
noise = (rng.standard_normal(length) * .08).clip(-1, 1).astype(np.float32)
tone = (.15 * np.sin(np.arange(length, dtype=np.float32) * np.float32(2 * np.pi * 440 / 16000))).astype(np.float32)
examples += [(f"noise-{length}", noise, None), (f"tone-{length}", tone, None)]
impulse = np.zeros(16000, dtype=np.float32)
impulse[0], impulse[7999], impulse[-1] = .9, -.9, .9
examples.append(("impulses-16000", impulse, None))
configuration = tf.compat.v1.ConfigProto(intra_op_parallelism_threads=4, inter_op_parallelism_threads=1, device_count={"GPU": 0})
records = []
with tf.Graph().as_default() as graph:
with tf.compat.v1.Session(graph=graph, config=configuration) as session:
tf.compat.v1.saved_model.loader.load(session, [], str(source))
inputs = graph.get_tensor_by_name(meta.signature_def["default"].inputs["default"].name)
endpoint = graph.get_tensor_by_name(meta.signature_def["default"].outputs["default"].name)
logmel = graph.get_tensor_by_name("Squeeze:0")
for name, wave, cached in examples:
features, mel_values = session.run([endpoint, logmel], {inputs: np.pad(wave, (11120, 0))[None]})
features = features[0, :, 0, :]
if features.shape != (len(wave) // 1280, 96) or not np.isfinite(features).all():
raise ValueError("Reference feature shape/finite failure")
arrays = {"waveform": wave, "features": features, "log_mel": mel_values[0]}
if cached is not None:
arrays["cached_source_features"] = cached
np.savez(output / (name + ".npz"), **arrays)
records.append({"name": name, "samples": len(wave), "frames": len(features),
"cached_source_max_abs": float(np.max(np.abs(cached - features))) if cached is not None else None})
report = {"stage": "reference_captured", "tensorflow": tf.__version__, "numpy": np.__version__, "source": identity,
"role": "32 metadata-selected dev-select query crops plus deterministic engineering stress fixtures",
"examples": records, "variables": len(weights), "window_shape": list(window.shape), "mel_shape": list(mel.shape),
"variable_npz_sha256": hashlib.sha256((output / "variables.npz").read_bytes()).hexdigest(),
"frontend_npz_sha256": hashlib.sha256((output / "frontend.npz").read_bytes()).hexdigest(),
"model_fitting": False, "final_evaluation": False}
(output / "report.json").write_text(json.dumps(report, indent=2) + "\n")
print(json.dumps({k: v for k, v in report.items() if k not in ("examples", "source")}, indent=2))
if __name__ == "__main__":
main()
|