saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-06 04:28:32 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Petermca/whisper dedup (#15365)

Browse source 
### Description
Apply `get_shared_initializers()` to the encoder and decoder subgraphs
of Whisper before chaining and exporting the full, final model.


### Motivation and Context
The Whisper export process has some overlap between the encoder and
decoder subgraphs due to the format of the BeamSearch contrib op.
Consequently, there is some shared model data that is duplicated in the
final exported product, which can result in a file size increase of
~40%. This PR takes the methods in `convert_generation.py` and applies
them during the whisper export process.

---------

Co-authored-by: Peter McAughan <petermca@microsoft.com>
This commit is contained in:
petermcaughan 2023-04-06 13:27:05 -07:00 committed by GitHub
parent dc1845a9c8
commit 2bd8e4a130
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 16 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
onnxruntime/python/tools/transformers/models/whisper/whisper_chain.py
View file

@ -1,7 +1,13 @@
import os
import sys
import onnx
from onnx import TensorProto, helper
from transformers import WhisperConfig
sys.path.append(os.path.join(os.path.dirname(__file__), "..", ".."))
from convert_generation import get_shared_initializers # noqa: E402
def add_attention_mask(model):
# Add attention mask - required by BeamSearch but unused in Pytorch
@ -13,11 +19,11 @@ def add_attention_mask(model):
def chain_model(args):
# Load encoder/decoder and insert necessary (but unused) graph inputs expected by BeamSearch op
encoder_model = onnx.load(args.encoder_path, load_external_data=True)
encoder_model = onnx.load_model(args.encoder_path, load_external_data=True)
encoder_model.graph.name = "encoderdecoderinit subgraph"
add_attention_mask(encoder_model)
decoder_model = onnx.load(args.decoder_path, load_external_data=True)
decoder_model = onnx.load_model(args.decoder_path, load_external_data=True)
decoder_model.graph.name = "decoder subgraph"
add_attention_mask(decoder_model)
@ -47,8 +53,6 @@ def chain_model(args):
helper.make_attribute("no_repeat_ngram_size", args.no_repeat_ngram_size),
helper.make_attribute("early_stopping", True),
helper.make_attribute("model_type", 2),
helper.make_attribute("decoder", decoder_model.graph),
helper.make_attribute("encoder", encoder_model.graph),
]
)
@ -84,7 +88,14 @@ def chain_model(args):
graph_outputs = [sequences]
# Initializers/opsets
initializers = []
# Delete shared data between decoder/encoder and move to larger graph initializers
initializers = get_shared_initializers(encoder_model, decoder_model)
node.attribute.extend(
[
helper.make_attribute("decoder", decoder_model.graph),
helper.make_attribute("encoder", encoder_model.graph),
]
)
opset_import = [helper.make_opsetid(domain="com.microsoft", version=1), helper.make_opsetid(domain="", version=17)]
beam_graph = helper.make_graph([node], "beam-search-test", graph_inputs, graph_outputs, initializers)