Add tf2pytorch wrapper in transformers tool (#5316)

* init checkin

* format

* refactor

* review comments

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()

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()."
    ]
   },
   {

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