From 2bd8e4a130bbe4e9a5cb63aec60b48f3b7346626 Mon Sep 17 00:00:00 2001 From: petermcaughan Date: Thu, 6 Apr 2023 13:27:05 -0700 Subject: [PATCH] Petermca/whisper dedup (#15365) ### 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 --- .../models/whisper/whisper_chain.py | 21 ++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/onnxruntime/python/tools/transformers/models/whisper/whisper_chain.py b/onnxruntime/python/tools/transformers/models/whisper/whisper_chain.py index c60260e3ac..bbdafdcb3e 100644 --- a/onnxruntime/python/tools/transformers/models/whisper/whisper_chain.py +++ b/onnxruntime/python/tools/transformers/models/whisper/whisper_chain.py @@ -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)