File size: 7,621 Bytes
5951253 | 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 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 | #!/usr/bin/env python3
"""Exports vinai/PhoWhisper-<size> to the ONNX files mobile apps run with ONNX Runtime (Android and iOS).
Writes out/<size>/repo/, laid out as the Hugging Face repo that gets published:
onnx/encoder_model.onnx fp32 encoder: input_features [1, n_mels, 3000] -> last_hidden_state
onnx/decoder_model_merged.onnx fp32 decoder with KV cache (use_cache_branch selects the first step)
onnx/encoder_model_quantized.onnx int8 (dynamic, MatMul weights) versions of the two above:
onnx/decoder_model_merged_quantized.onnx the files mobile apps should download
config.json, generation_config.json, preprocessor_config.json, tokenizer files
Weights over 2 GB (medium, large) go to a sidecar <name>.onnx_data next to the graph.
Usage:
.venv/bin/python -I export.py tiny|base|small|medium|large
"""
import shutil
import sys
from pathlib import Path
import numpy as np
import onnx
from onnx import helper, numpy_helper
from onnxruntime.quantization import QuantType, quantize_dynamic
HERE = Path(__file__).resolve().parent
SIZES = ("tiny", "base", "small", "medium", "large")
KEEP = ("encoder_model", "decoder_model_merged")
COPIED = (
"config.json",
"generation_config.json",
"preprocessor_config.json",
"tokenizer.json",
"tokenizer_config.json",
"vocab.json",
"merges.txt",
"added_tokens.json",
"special_tokens_map.json",
"normalizer.json",
)
# Protobuf messages cannot exceed 2 GB; above this the weights go to a sidecar file.
EXTERNAL_DATA_THRESHOLD = int(1.8 * 1024**3)
def check_and_save_large(original):
"""Wraps Optimum's check_and_save_model so it also saves models over 2 GB (PhoWhisper-large).
It starts with model.ByteSize(), which current protobuf refuses to compute over 2 GB, so a
large model goes straight to an external-data save here instead.
"""
def check_and_save(model: onnx.ModelProto, save_path) -> None:
if sum(len(t.raw_data) for t in model.graph.initializer) < EXTERNAL_DATA_THRESHOLD:
return original(model, save_path)
path = Path(save_path)
save(model, path)
onnx.checker.check_model(str(path))
return check_and_save
def export(size: str, raw: Path) -> None:
"""Runs the Optimum export (encoder, decoder, decoder with past, merged decoder) into raw/."""
if (raw / "decoder_model_merged.onnx").exists():
print(f"reusing {raw}")
return
shutil.rmtree(raw, ignore_errors=True)
from optimum.exporters.onnx import convert, main_export
from optimum.onnx import graph_transformations
patched = check_and_save_large(graph_transformations.check_and_save_model)
graph_transformations.check_and_save_model = patched
convert.check_and_save_model = patched
main_export(
f"vinai/PhoWhisper-{size}",
output=raw,
task="automatic-speech-recognition-with-past",
opset=17,
)
def save(model: onnx.ModelProto, path: Path) -> None:
for stale in (path, path.with_name(path.name + "_data")):
stale.unlink(missing_ok=True)
weights = sum(len(t.raw_data) for t in model.graph.initializer)
if weights > EXTERNAL_DATA_THRESHOLD:
onnx.save_model(
model, str(path),
save_as_external_data=True, all_tensors_to_one_file=True,
location=path.name + "_data", size_threshold=1024,
)
else:
onnx.save_model(model, str(path))
def copy_fp32(raw: Path, onnx_dir: Path) -> None:
for name in KEEP:
# Re-saving gathers the weights into one sidecar (Optimum writes one file per tensor for large models).
model = onnx.load(str(raw / f"{name}.onnx"), load_external_data=True)
save(model, onnx_dir / f"{name}.onnx")
print(f"wrote {onnx_dir / name}.onnx")
def rename_input(graph: onnx.GraphProto, old: str, new: str) -> None:
for node in graph.node:
node.input[:] = [new if i == old else i for i in node.input]
for attr in node.attribute:
if attr.type == onnx.AttributeProto.GRAPH:
rename_input(attr.g, old, new)
def pretranspose_constants(graph: onnx.GraphProto, initializers: dict[str, onnx.TensorProto]) -> int:
"""Replaces Transpose(initializer) with a transposed initializer, in graph and its subgraphs.
The decoder's LM head shares the token embedding and multiplies by its transpose, computed at run
time. The quantizer skips a MatMul whose weight comes out of a node, which would leave the LM head
fp32 (80 MB for tiny), so the transpose is folded into a constant first.
"""
folded = 0
for node in list(graph.node):
for attr in node.attribute:
if attr.type == onnx.AttributeProto.GRAPH:
folded += pretranspose_constants(attr.g, initializers)
if node.op_type != "Transpose" or node.input[0] not in initializers:
continue
perm = next((list(a.ints) for a in node.attribute if a.name == "perm"), None)
folded_name = node.input[0] + "_transposed"
if folded_name not in initializers:
weight = numpy_helper.to_array(initializers[node.input[0]])
initializers[folded_name] = numpy_helper.from_array(np.ascontiguousarray(weight.transpose(perm)), folded_name)
graph.node.remove(node)
rename_input(graph, node.output[0], folded_name)
folded += 1
return folded
def quantize(onnx_dir: Path, tmp: Path) -> None:
tmp.mkdir(exist_ok=True)
for name in KEEP:
model = onnx.load(str(onnx_dir / f"{name}.onnx"), load_external_data=True)
initializers = {t.name: t for t in model.graph.initializer}
known = set(initializers)
if pretranspose_constants(model.graph, initializers):
model.graph.initializer.extend(t for n, t in initializers.items() if n not in known)
src = tmp / f"{name}.onnx"
save(model, src)
dst = onnx_dir / f"{name}_quantized.onnx"
# MatMul and Gather (the token embedding), also inside the merged decoder's If branches.
# The encoder's two Conv layers stay fp32, since ConvInteger has no kernel on
# some mobile builds. Signed weights, no reduce_range: ARM's int8 dot product kernels need neither.
quantize_dynamic(
src, dst,
op_types_to_quantize=["MatMul", "Gather"],
weight_type=QuantType.QInt8,
per_channel=False,
reduce_range=False,
use_external_data_format=src.with_name(src.name + "_data").exists(),
extra_options={"MatMulConstBOnly": True, "EnableSubgraph": True},
)
model = onnx.load(str(dst), load_external_data=True)
save(model, dst) # int8 weights of every size fit under 2 GB, so this inlines them again
dst.with_name(dst.name + ".data").unlink(missing_ok=True) # quantize_dynamic's sidecar, now inlined
print(f"wrote {dst}")
shutil.rmtree(tmp)
def copy_tokenizer(raw: Path, repo: Path) -> None:
for name in COPIED:
if (raw / name).exists():
shutil.copy2(raw / name, repo / name)
def main() -> None:
if len(sys.argv) != 2 or sys.argv[1] not in SIZES:
sys.exit(f"usage: {sys.argv[0]} {'|'.join(SIZES)}")
size = sys.argv[1]
work = HERE / "out" / size
raw, repo = work / "raw", work / "repo"
onnx_dir = repo / "onnx"
onnx_dir.mkdir(parents=True, exist_ok=True)
export(size, raw)
copy_fp32(raw, onnx_dir)
quantize(onnx_dir, work / "tmp")
copy_tokenizer(raw, repo)
if __name__ == "__main__":
main()
|