From caed6c264c869dcffd181a5e9628ae1962c914a4 Mon Sep 17 00:00:00 2001 From: Ye Wang <52801275+wangyems@users.noreply.github.com> Date: Thu, 1 Oct 2020 13:58:58 -0700 Subject: [PATCH] Add tf2pytorch wrapper in transformers tool (#5316) * init checkin * format * refactor * review comments --- .../convert_tf_models_to_pytorch.py | 159 ++++++++++++++++++ ...low_Keras_Bert-Squad_OnnxRuntime_CPU.ipynb | 4 +- .../tools/transformers/onnx_exporter.py | 39 ++++- 3 files changed, 193 insertions(+), 9 deletions(-) create mode 100644 onnxruntime/python/tools/transformers/convert_tf_models_to_pytorch.py diff --git a/onnxruntime/python/tools/transformers/convert_tf_models_to_pytorch.py b/onnxruntime/python/tools/transformers/convert_tf_models_to_pytorch.py new file mode 100644 index 0000000000..30ec2e2479 --- /dev/null +++ b/onnxruntime/python/tools/transformers/convert_tf_models_to_pytorch.py @@ -0,0 +1,159 @@ +#------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +#-------------------------------------------------------------------------- + +import glob +import os +import requests + +TFMODELS = { + "bert-base-uncased": + ("bert", "BertConfig", "", "https://storage.googleapis.com/bert_models/2018_10_18/uncased_L-12_H-768_A-12.zip"), + "bert-base-cased": + ("bert", "BertConfig", "", "https://storage.googleapis.com/bert_models/2019_05_30/wwm_cased_L-24_H-1024_A-16.zip"), + "bert-large-uncased": + ("bert", "BertConfig", "", "https://storage.googleapis.com/bert_models/2018_10_18/uncased_L-24_H-1024_A-16.zip"), + "albert-base": ("albert", "AlbertConfig", "", "https://storage.googleapis.com/albert_models/albert_base_v1.tar.gz"), + "albert-large": + ("albert", "AlbertConfig", "", "https://storage.googleapis.com/albert_models/albert_large_v1.tar.gz"), + "gpt-2-117M": ("gpt2", "GPT2Config", "GPT2Model", "https://storage.googleapis.com/gpt-2/models/117M"), + "gpt-2-124M": ("gpt2", "GPT2Config", "GPT2Model", "https://storage.googleapis.com/gpt-2/models/124M") +} + + +def download_compressed_file(tf_ckpt_url, ckpt_dir): + r = requests.get(tf_ckpt_url) + compressed_file_name = tf_ckpt_url.split("/")[-1] + compressed_file_dir = os.path.join(ckpt_dir, compressed_file_name) + with open(compressed_file_dir, 'wb') as f: + f.write(r.content) + return compressed_file_dir + + +def get_ckpt_prefix_path(ckpt_dir): + # get prefix + sub_folder_dir = None + for o in os.listdir(ckpt_dir): + sub_folder_dir = os.path.join(ckpt_dir, o) + break + if os.path.isfile(sub_folder_dir): + sub_folder_dir = ckpt_dir + unique_file_name = str(glob.glob(sub_folder_dir + "/*data-00000-of-00001")) + prefix = (unique_file_name.rpartition('.')[0]).split("/")[-1] + + return os.path.join(sub_folder_dir, prefix) + + +def download_tf_checkpoint(model_name, tf_models_dir="tf_models"): + import pathlib + base_dir = os.path.join(pathlib.Path(__file__).parent.absolute(), tf_models_dir) + ckpt_dir = os.path.join(base_dir, model_name) + + if not os.path.exists(ckpt_dir): + os.makedirs(ckpt_dir) + + tf_ckpt_url = TFMODELS[model_name][3] + + import re + if (re.search('.zip$', tf_ckpt_url) != None): + zip_dir = download_compressed_file(tf_ckpt_url, ckpt_dir) + + # unzip file + import zipfile + with zipfile.ZipFile(zip_dir, 'r') as zip_ref: + zip_ref.extractall(ckpt_dir) + os.remove(zip_dir) + + return get_ckpt_prefix_path(ckpt_dir) + + elif (re.search('.tar.gz$', tf_ckpt_url) != None): + tar_dir = download_compressed_file(tf_ckpt_url, ckpt_dir) + + # untar file + import tarfile + with tarfile.open(tar_dir, 'r') as tar_ref: + tar_ref.extractall(ckpt_dir) + os.remove(tar_dir) + + return get_ckpt_prefix_path(ckpt_dir) + + else: + for filename in ['checkpoint', 'model.ckpt.data-00000-of-00001', 'model.ckpt.index', 'model.ckpt.meta']: + r = requests.get(tf_ckpt_url + "/" + filename) + with open(os.path.join(ckpt_dir, filename), 'wb') as f: + f.write(r.content) + + return get_ckpt_prefix_path(ckpt_dir) + + +def init_pytorch_model(model_name, tf_checkpoint_path): + config_name = TFMODELS[model_name][1] + config_module = __import__("transformers", fromlist=[config_name]) + model_config = getattr(config_module, config_name) + + parent_path = tf_checkpoint_path.rpartition('/')[0] + config_path = glob.glob(parent_path + "/*config.json") + config = model_config() if len(config_path) is 0 else model_config.from_json_file(str(config_path[0])) + + if TFMODELS[model_name][2] is "": + from transformers import AutoModelForPreTraining + init_model = AutoModelForPreTraining.from_config(config) + else: + model_categroy_name = TFMODELS[model_name][2] + module = __import__("transformers", fromlist=[model_categroy_name]) + model_categroy = getattr(module, model_categroy_name) + init_model = model_categroy(config) + return config, init_model + + +def convert_tf_checkpoint_to_pytorch(model_name, config, init_model, tf_checkpoint_path, is_tf2): + load_tf_weight_func_name = "load_tf_weights_in_" + TFMODELS[model_name][0] + + module = __import__("transformers", fromlist=[load_tf_weight_func_name]) + + if is_tf2 is False: + load_tf_weight_func = getattr(module, load_tf_weight_func_name) + else: + if TFMODELS[model_name][0] is not "bert": + raise NotImplementedError("Only support tf2 ckeckpoint for Bert model") + from transformers import convert_bert_original_tf2_checkpoint_to_pytorch + load_tf_weight_func = convert_bert_original_tf2_checkpoint_to_pytorch.load_tf2_weights_in_bert + + # Expect transformers team will unify the order of signature in the future + model = load_tf_weight_func(init_model, config, tf_checkpoint_path) if is_tf2 is False else load_tf_weight_func( + init_model, tf_checkpoint_path, config) + model.eval() + return model + + +def tf2pt_pipeline(model_name, is_tf2=False): + if model_name not in TFMODELS: + raise NotImplementedError(model_name + " not implemented") + tf_checkpoint_path = download_tf_checkpoint(model_name) + config, init_model = init_pytorch_model(model_name, tf_checkpoint_path) + model = convert_tf_checkpoint_to_pytorch(model_name, config, init_model, tf_checkpoint_path, is_tf2) + # Could then use the model in Benchmark + return config, model + +def tf2pt_pipeline_test(): + # For test on linux only + import logging + import torch + logger = logging.getLogger('') + for model_name in TFMODELS.keys(): + config, model = tf2pt_pipeline(model_name) + assert(config.model_type is TFMODELS[model_name][0]) + + input = torch.randint(low=0, + high=config.vocab_size - 1, + size=(4, 128), + dtype=torch.long) + try: + model(input) + except RuntimeError as e: + logger.exception(e) + + +if __name__ == '__main__': + tf2pt_pipeline_test() diff --git a/onnxruntime/python/tools/transformers/notebooks/Tensorflow_Keras_Bert-Squad_OnnxRuntime_CPU.ipynb b/onnxruntime/python/tools/transformers/notebooks/Tensorflow_Keras_Bert-Squad_OnnxRuntime_CPU.ipynb index caba3614fa..c2067f5325 100644 --- a/onnxruntime/python/tools/transformers/notebooks/Tensorflow_Keras_Bert-Squad_OnnxRuntime_CPU.ipynb +++ b/onnxruntime/python/tools/transformers/notebooks/Tensorflow_Keras_Bert-Squad_OnnxRuntime_CPU.ipynb @@ -213,7 +213,9 @@ "source": [ "## 3. Export model to ONNX using Keras2onnx\n", "\n", - "Now we use Keras2onnx to export the model to ONNX format. It takes about 3 minutes for bert-base, or 18 minutes for bert-large model." + "Now we use Keras2onnx to export the model to ONNX format. It takes about 3 minutes for bert-base, or 18 minutes for bert-large model.", + "\n", + "Note that we could also convert tensorflow checkpoints to pytorch(supported by huggingface team, ref:https://huggingface.co/transformers/converting_tensorflow_models.html) and then convert to onnx using torch.onnx.export()." ] }, { diff --git a/onnxruntime/python/tools/transformers/onnx_exporter.py b/onnxruntime/python/tools/transformers/onnx_exporter.py index 6f9ad1fc41..0dcc3c4b6c 100644 --- a/onnxruntime/python/tools/transformers/onnx_exporter.py +++ b/onnxruntime/python/tools/transformers/onnx_exporter.py @@ -231,6 +231,34 @@ def load_pretrained_model(model_name, config, cache_dir, custom_model_class, is_ return model_class.from_pretrained(model_name, config=config, cache_dir=cache_dir, use_cdn=use_cdn) +def load_pt_model(model_name, model_class, cache_dir): + config = AutoConfig.from_pretrained(model_name, cache_dir=cache_dir) + if hasattr(config, 'return_dict'): + config.return_dict = False + + model = load_pretrained_model(model_name, config=config, cache_dir=cache_dir, custom_model_class=model_class) + + return config, model + + +def load_tf_model(model_name, model_class, cache_dir): + config = AutoConfig.from_pretrained(model_name, cache_dir=cache_dir) + + model = load_pretrained_model(model_name, config=config, cache_dir=cache_dir, custom_model_class=model_class, is_tf_model=True) + + return config, model + + +# For test only +def load_pt_model_from_tf(model_name): + # Note that we could get pt model from tf, but model source and its structure in this case is different from directly using + # load_pt_model() and load_tf_model() even with the same name. Therefore it should not be used for comparing with them + from convert_tf_models_to_pytorch import tf2pt_pipeline + config, model = tf2pt_pipeline(model_name) + + return config, model + + def validate_and_optimize_onnx(model_name, use_external_data_format, model_type, onnx_dir, input_names, use_gpu, precision, optimize_onnx, validate_onnx, use_raw_attention_mask, overwrite, config, model_fusion_statistics, onnx_model_path, example_inputs, example_outputs_flatten): @@ -269,11 +297,8 @@ def export_onnx_model_from_pt(model_name, opset_version, use_external_data_forma input_names, use_gpu, precision, optimize_onnx, validate_onnx, use_raw_attention_mask, overwrite, model_fusion_statistics): - config = AutoConfig.from_pretrained(model_name, cache_dir=cache_dir) - if hasattr(config, 'return_dict'): - config.return_dict = False - - model = load_pretrained_model(model_name, config=config, cache_dir=cache_dir, custom_model_class=model_class) + config, model = load_pt_model(model_name, model_class, cache_dir) + # config, model = load_pt_model_from_tf(model_name) model.cpu() tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir) @@ -326,9 +351,7 @@ def export_onnx_model_from_tf(model_name, opset_version, use_external_data_forma input_names, use_gpu, precision, optimize_onnx, validate_onnx, use_raw_attention_mask, overwrite, model_fusion_statistics): - config = AutoConfig.from_pretrained(model_name, cache_dir=cache_dir) - - model = load_pretrained_model(model_name, config=config, cache_dir=cache_dir, custom_model_class=model_class, is_tf_model=True) + config, model = load_tf_model(model_name, model_class, cache_dir) model._saved_model_inputs_spec = None