Conversion script to export transformers models to ONNX IR. (#4253)
* Added generic ONNX conversion script for PyTorch model. * WIP initial TF support. * TensorFlow/Keras ONNX export working. * Print framework version info * Add possibility to check the model is correctly loading on ONNX runtime. * Remove quantization option. * Specify ONNX opset version when exporting. * Formatting. * Remove unused imports. * Make functions more generally reusable from other part of the code. * isort happy. * flake happy * Export only feature-extraction for now * Correctly check inputs order / filter before export. * Removed task variable * Fix invalid args call in load_graph_from_args. * Fix invalid args call in convert. * Fix invalid args call in infer_shapes. * Raise exception and catch in caller function instead of exit. * Add 04-onnx-export.ipynb notebook * More WIP on the notebook * Remove unused imports * Simplify & remove unused constants. * Export with constant_folding in PyTorch * Let's try to put function args in the right order this time ... * Disable external_data_format temporary * ONNX notebook draft ready. * Updated notebooks charts + wording * Correct error while exporting last chart in notebook. * Adressing @LysandreJik comment. * Set ONNX opset to 11 as default value. * Set opset param mandatory * Added ONNX export unittests * Quality. * flake8 happy * Add keras2onnx dependency on extras["tf"] * Pin keras2onnx on github master to v1.6.5 * Second attempt. * Third attempt. * Use the right repo URL this time ... * Do the same for onnxconverter-common * Added keras2onnx and onnxconveter-common to 1.7.0 to supports TF2.2 * Correct commit hash. * Addressing PR review: Optimization are enabled by default. * Addressing PR review: small changes in the notebook * setup.py comment about keras2onnx versioning.
This commit is contained in:
Funtowicz Morgan authored and GitHub committed
1 parent
2d05480174
commit
db0076a9df
5 files changed
+804
-3
No files matched your search
@@ -0,0 +1,462 @@
|
|||||||
|
{
|
||||||
|
"cells": [
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {
|
||||||
|
"colab_type": "text",
|
||||||
|
"id": "jBasof3bv1LB"
|
||||||
|
},
|
||||||
|
"source": [
|
||||||
|
"<h1><center>How to export 🤗 Transformers Models to ONNX ?<h1><center>"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"[ONNX](http://onnx.ai/) is open format for machine learning models. It allows to save your neural network's computation graph in a framework agnostic way, which might be particulary helpful when deploying deep learning models.\n",
|
||||||
|
"\n",
|
||||||
|
"Indeed, businesses might have other requirements _(languages, hardware, ...)_ for which the training framework might not be the best suited in inference scenarios. In that context, having a representation of the actual computation graph that can be shared accross various business units and logics across an organization might be a desirable component.\n",
|
||||||
|
"\n",
|
||||||
|
"Along with the serialization format, ONNX also provides a runtime library which allows efficient and hardware specific execution of the ONNX graph. This is done through the [onnxruntime](https://microsoft.github.io/onnxruntime/) project and already includes collaborations with many hardware vendors to seamlessly deploy models on various platforms.\n",
|
||||||
|
"\n",
|
||||||
|
"Through this notebook we'll walk you through the process to convert a PyTorch or TensorFlow transformers model to the [ONNX](http://onnx.ai/) and leverage [onnxruntime](https://microsoft.github.io/onnxruntime/) to run inference tasks on models from 🤗 __transformers__"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {
|
||||||
|
"colab_type": "text",
|
||||||
|
"id": "yNnbrSg-5e1s"
|
||||||
|
},
|
||||||
|
"source": [
|
||||||
|
"## Exporting 🤗 transformers model to ONNX\n",
|
||||||
|
"\n",
|
||||||
|
"---\n",
|
||||||
|
"\n",
|
||||||
|
"Exporting models _(either PyTorch or TensorFlow)_ is easily achieved through the conversion tool provided as part of 🤗 __transformers__ repository. \n",
|
||||||
|
"\n",
|
||||||
|
"Under the hood the process is sensibly the following: \n",
|
||||||
|
"\n",
|
||||||
|
"1. Allocate the model from transformers (**PyTorch or TensorFlow**)\n",
|
||||||
|
"2. Forward dummy inputs through the model this way **ONNX** can record the set of operations executed\n",
|
||||||
|
"3. Optionally define dynamic axes on input and output tensors\n",
|
||||||
|
"4. Save the graph along with the network parameters"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {
|
||||||
|
"scrolled": false
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"!pip install --upgrade git+https://github.com/huggingface/transformers@onnx-export"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {
|
||||||
|
"colab": {},
|
||||||
|
"colab_type": "code",
|
||||||
|
"id": "PwAaOchY4N2-"
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"!rm -rf onnx/\n",
|
||||||
|
"from transformers.convert_graph_to_onnx import convert\n",
|
||||||
|
"\n",
|
||||||
|
"# Handles all the above steps for you\n",
|
||||||
|
"convert(framework=\"pt\", model=\"bert-base-cased\", output=\"onnx/bert-base-cased.onnx\", opset=11)\n",
|
||||||
|
"\n",
|
||||||
|
"# Tensorflow \n",
|
||||||
|
"# convert(framework=\"tf\", model=\"bert-base-cased\", output=\"onnx/bert-base-cased.onnx\", opset=11)"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"## How to leverage runtime for inference over an ONNX graph\n",
|
||||||
|
"\n",
|
||||||
|
"---\n",
|
||||||
|
"\n",
|
||||||
|
"As mentionned in the introduction, **ONNX** is a serialization format and many side projects can load the saved graph and run the actual computations from it. Here, we'll focus on the official [onnxruntime](https://microsoft.github.io/onnxruntime/). The runtime is implemented in C++ for performance reasons and provides API/Bindings for C++, C, C#, Java and Python.\n",
|
||||||
|
"\n",
|
||||||
|
"In the case of this notebook, we will use the Python API to highlight how to load a serialized **ONNX** graph and run inference workload on various backends through **onnxruntime**.\n",
|
||||||
|
"\n",
|
||||||
|
"**onnxruntime** is available on pypi:\n",
|
||||||
|
"\n",
|
||||||
|
"- onnxruntime: ONNX + MLAS (Microsoft Linear Algebra Subprograms)\n",
|
||||||
|
"- onnxruntime-gpu: ONNX + MLAS + CUDA\n"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {
|
||||||
|
"pycharm": {
|
||||||
|
"name": "#%%\n"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"!pip install transformers onnxruntime-gpu onnx psutil matplotlib"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {
|
||||||
|
"colab_type": "text",
|
||||||
|
"id": "-gP08tHfBvgY"
|
||||||
|
},
|
||||||
|
"source": [
|
||||||
|
"## Preparing for an Inference Session\n",
|
||||||
|
"\n",
|
||||||
|
"---\n",
|
||||||
|
"\n",
|
||||||
|
"Inference is done using a specific backend definition which turns on hardware specific optimizations of the graph. \n",
|
||||||
|
"\n",
|
||||||
|
"Optimizations are basically of three kinds: \n",
|
||||||
|
"\n",
|
||||||
|
"- **Constant Folding**: Convert static variables to constants in the graph \n",
|
||||||
|
"- **Deadcode Elimination**: Remove nodes never accessed in the graph\n",
|
||||||
|
"- **Operator Fusing**: Merge multiple instruction into one (Linear -> ReLU can be fused to be LinearReLU)\n",
|
||||||
|
"\n",
|
||||||
|
"All of this is done on **onnxruntime** by settings specific `SessionOptions`:"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 2,
|
||||||
|
"metadata": {
|
||||||
|
"pycharm": {
|
||||||
|
"name": "#%%\n"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"from os import environ\n",
|
||||||
|
"from psutil import cpu_count\n",
|
||||||
|
"\n",
|
||||||
|
"# Constants from the performance optimization available in onnxruntime\n",
|
||||||
|
"# It needs to be done before importing onnxruntime\n",
|
||||||
|
"environ[\"OMP_NUM_THREADS\"] = str(cpu_count(logical=True))\n",
|
||||||
|
"environ[\"OMP_WAIT_POLICY\"] = 'ACTIVE'\n",
|
||||||
|
"\n",
|
||||||
|
"from onnxruntime import InferenceSession, SessionOptions, get_all_providers"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 3,
|
||||||
|
"metadata": {
|
||||||
|
"colab": {},
|
||||||
|
"colab_type": "code",
|
||||||
|
"id": "2k-jHLfdcTFS"
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"def create_model_for_provider(model_path: str, provider: str) -> InferenceSession: \n",
|
||||||
|
" \n",
|
||||||
|
" assert provider in get_all_providers(), f\"provider {provider} not found, {get_all_providers()}\"\n",
|
||||||
|
"\n",
|
||||||
|
" # Few properties than might have an impact on performances (provided by MS)\n",
|
||||||
|
" options = SessionOptions()\n",
|
||||||
|
" options.intra_op_num_threads = 1\n",
|
||||||
|
"\n",
|
||||||
|
" # Load the model as a graph and prepare the CPU backend \n",
|
||||||
|
" return InferenceSession(model_path, options, providers=[provider])"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {
|
||||||
|
"colab_type": "text",
|
||||||
|
"id": "teJdG3amE-hR"
|
||||||
|
},
|
||||||
|
"source": [
|
||||||
|
"## Forwarding through our optimized ONNX model running on CPU\n",
|
||||||
|
"\n",
|
||||||
|
"---\n",
|
||||||
|
"\n",
|
||||||
|
"When the model is loaded for inference over a specific provider, for instance **CPUExecutionProvider** as above, an optimized graph can be saved. This graph will might include various optimizations, and you might be able to see some **higher-level** operations in the graph _(through [Netron](https://github.com/lutzroeder/Netron) for instance)_ such as:\n",
|
||||||
|
"- **EmbedLayerNormalization**\n",
|
||||||
|
"- **Attention**\n",
|
||||||
|
"- **FastGeLU**\n",
|
||||||
|
"\n",
|
||||||
|
"These operations are an example of the kind of optimization **onnxruntime** is doing, for instance here gathering multiple operations into bigger one _(Operator Fusing)_."
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 4,
|
||||||
|
"metadata": {
|
||||||
|
"colab": {
|
||||||
|
"base_uri": "https://localhost:8080/",
|
||||||
|
"height": 34
|
||||||
|
},
|
||||||
|
"colab_type": "code",
|
||||||
|
"id": "dmC22kJfVGYe",
|
||||||
|
"outputId": "f3aba5dc-15c0-4f82-b38c-1bbae1bf112e"
|
||||||
|
},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"name": "stdout",
|
||||||
|
"output_type": "stream",
|
||||||
|
"text": [
|
||||||
|
"Sequence output: (1, 6, 768), Pooled output: (1, 768)\n"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"from transformers import BertTokenizerFast\n",
|
||||||
|
"\n",
|
||||||
|
"tokenizer = BertTokenizerFast.from_pretrained(\"bert-base-cased\")\n",
|
||||||
|
"cpu_model = create_model_for_provider(\"onnx/bert-base-cased.onnx\", \"CPUExecutionProvider\")\n",
|
||||||
|
"\n",
|
||||||
|
"# Inputs are provided through numpy array\n",
|
||||||
|
"model_inputs = tokenizer.encode_plus(\"My name is Bert\", return_tensors=\"pt\")\n",
|
||||||
|
"inputs_onnx = {k: v.cpu().detach().numpy() for k, v in model_inputs.items()}\n",
|
||||||
|
"\n",
|
||||||
|
"# Run the model (None = get all the outputs)\n",
|
||||||
|
"sequence, pooled = cpu_model.run(None, inputs_onnx)\n",
|
||||||
|
"\n",
|
||||||
|
"# Print information about outputs\n",
|
||||||
|
"\n",
|
||||||
|
"print(f\"Sequence output: {sequence.shape}, Pooled output: {pooled.shape}\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {
|
||||||
|
"colab_type": "text",
|
||||||
|
"id": "Kda1e7TkEqNR"
|
||||||
|
},
|
||||||
|
"source": [
|
||||||
|
"## Benchmarking different CPU & GPU providers\n",
|
||||||
|
"\n",
|
||||||
|
"_**Disclamer: results may vary from the actual hardware used to run the model**_"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 5,
|
||||||
|
"metadata": {
|
||||||
|
"colab": {
|
||||||
|
"base_uri": "https://localhost:8080/",
|
||||||
|
"height": 170
|
||||||
|
},
|
||||||
|
"colab_type": "code",
|
||||||
|
"id": "WcdFZCvImVig",
|
||||||
|
"outputId": "bfd779a1-0bc7-42db-8587-e52a485ec5e3"
|
||||||
|
},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"name": "stdout",
|
||||||
|
"output_type": "stream",
|
||||||
|
"text": [
|
||||||
|
"Doing GPU inference on TITAN RTX\n"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"name": "stderr",
|
||||||
|
"output_type": "stream",
|
||||||
|
"text": [
|
||||||
|
"Warming up: 100%|██████████| 10/10 [00:00<00:00, 333.82it/s]\n",
|
||||||
|
"Tracking inference time on CUDAExecutionProvider: 100%|██████████| 100/100 [00:00<00:00, 521.76it/s]\n",
|
||||||
|
"Warming up: 100%|██████████| 10/10 [00:00<00:00, 62.95it/s]\n",
|
||||||
|
"Tracking inference time on CPUExecutionProvider: 100%|██████████| 100/100 [00:01<00:00, 68.65it/s]\n",
|
||||||
|
"Warming up: 100%|██████████| 10/10 [00:00<00:00, 69.72it/s]\n",
|
||||||
|
"Tracking inference time on TensorrtExecutionProvider: 100%|██████████| 100/100 [00:01<00:00, 71.31it/s]\n",
|
||||||
|
"Warming up: 100%|██████████| 10/10 [00:00<00:00, 66.28it/s]\n",
|
||||||
|
"Tracking inference time on DnnlExecutionProvider: 100%|██████████| 100/100 [00:01<00:00, 72.03it/s]\n"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"from torch.cuda import get_device_name\n",
|
||||||
|
"from contextlib import contextmanager\n",
|
||||||
|
"from dataclasses import dataclass\n",
|
||||||
|
"from time import time\n",
|
||||||
|
"from tqdm import trange\n",
|
||||||
|
"\n",
|
||||||
|
"print(f\"Doing GPU inference on {get_device_name(0)}\", flush=True)\n",
|
||||||
|
"\n",
|
||||||
|
"@contextmanager\n",
|
||||||
|
"def track_infer_time(buffer: [int]):\n",
|
||||||
|
" start = time()\n",
|
||||||
|
" yield\n",
|
||||||
|
" end = time()\n",
|
||||||
|
"\n",
|
||||||
|
" buffer.append(end - start)\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"@dataclass\n",
|
||||||
|
"class OnnxInferenceResult:\n",
|
||||||
|
" model_inference_time: [int] \n",
|
||||||
|
" optimized_model_path: str\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"# All the providers we'll be using in the test\n",
|
||||||
|
"results = {}\n",
|
||||||
|
"providers = [\n",
|
||||||
|
" \"CUDAExecutionProvider\",\n",
|
||||||
|
" \"CPUExecutionProvider\", \n",
|
||||||
|
" \"TensorrtExecutionProvider\",\n",
|
||||||
|
" \"DnnlExecutionProvider\", \n",
|
||||||
|
"]\n",
|
||||||
|
"\n",
|
||||||
|
"# Iterate over all the providers\n",
|
||||||
|
"for provider in providers:\n",
|
||||||
|
"\n",
|
||||||
|
" # Create the model with the specified provider\n",
|
||||||
|
" model = create_model_for_provider(\"onnx/bert-base-cased.onnx\", provider)\n",
|
||||||
|
"\n",
|
||||||
|
" # Keep track of the inference time\n",
|
||||||
|
" time_buffer = []\n",
|
||||||
|
"\n",
|
||||||
|
" # Warm up the model\n",
|
||||||
|
" for _ in trange(10, desc=\"Warming up\"):\n",
|
||||||
|
" model.run(None, inputs_onnx)\n",
|
||||||
|
"\n",
|
||||||
|
" # Compute \n",
|
||||||
|
" for _ in trange(100, desc=f\"Tracking inference time on {provider}\"):\n",
|
||||||
|
" with track_infer_time(time_buffer):\n",
|
||||||
|
" model.run(None, inputs_onnx)\n",
|
||||||
|
"\n",
|
||||||
|
" # Store the result\n",
|
||||||
|
" results[provider] = OnnxInferenceResult(\n",
|
||||||
|
" time_buffer,\n",
|
||||||
|
" model.get_session_options().optimized_model_filepath\n",
|
||||||
|
" )"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 7,
|
||||||
|
"metadata": {
|
||||||
|
"colab": {
|
||||||
|
"base_uri": "https://localhost:8080/",
|
||||||
|
"height": 51
|
||||||
|
},
|
||||||
|
"colab_type": "code",
|
||||||
|
"id": "PS_49goe197g",
|
||||||
|
"outputId": "0ef0f70c-f5a7-46a0-949a-1a93f231d193"
|
||||||
|
},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"name": "stderr",
|
||||||
|
"output_type": "stream",
|
||||||
|
"text": [
|
||||||
|
"Warming up: 100%|██████████| 10/10 [00:00<00:00, 18.04it/s]\n",
|
||||||
|
"Tracking inference time on PyTorch: 100%|██████████| 100/100 [00:05<00:00, 18.88it/s]\n"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"from transformers import BertModel\n",
|
||||||
|
"\n",
|
||||||
|
"# Add PyTorch to the providers\n",
|
||||||
|
"model_pt = BertModel.from_pretrained(\"bert-base-cased\")\n",
|
||||||
|
"for _ in trange(10, desc=\"Warming up\"):\n",
|
||||||
|
" model_pt(**model_inputs)\n",
|
||||||
|
"\n",
|
||||||
|
"# Compute \n",
|
||||||
|
"time_buffer = []\n",
|
||||||
|
"for _ in trange(100, desc=f\"Tracking inference time on PyTorch\"):\n",
|
||||||
|
" with track_infer_time(time_buffer):\n",
|
||||||
|
" model_pt(**model_inputs)\n",
|
||||||
|
"\n",
|
||||||
|
"# Store the result\n",
|
||||||
|
"results[\"Pytorch\"] = OnnxInferenceResult(\n",
|
||||||
|
" time_buffer, \n",
|
||||||
|
" model.get_session_options().optimized_model_filepath\n",
|
||||||
|
") "
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"## Show the inference performance of each providers \n",
|
||||||
|
"\n",
|
||||||
|
"_Note: PyTorch model benchmark is run on CPU_"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 24,
|
||||||
|
"metadata": {
|
||||||
|
"colab": {
|
||||||
|
"base_uri": "https://localhost:8080/",
|
||||||
|
"height": 676
|
||||||
|
},
|
||||||
|
"colab_type": "code",
|
||||||
|
"id": "dj-rS8AcqRZQ",
|
||||||
|
"outputId": "b4bf07d1-a7b4-4eff-e6bd-d5d424fd17fb"
|
||||||
|
},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"data": {
|
||||||
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAABRoAAAPeCAYAAABjjKazAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAAPYQAAD2EBqD+naQAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjEsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+j8jraAAAgAElEQVR4nOzdd5SU5d34/88sZZe2gIBABAHBgjWICYIKqFhBJKKAGoOFkMRYwBgjMRGJKIpGUaOiMQ+xrAELEDSPgkY0FuyxRKOiAeVBBUQBqSI73z/87fwYdqnX6oJ5vc7Zc3avuWfua+4pR97eJZPNZrMBAAAAAJCgoKonAAAAAABs+4RGAAAAACCZ0AgAAAAAJBMaAQAAAIBkQiMAAAAAkExoBAAAAACSCY0AAAAAQDKhEQAAAABIJjQCAAAAAMmERgBgm9K6des49dRTt/j+V111Vey0005RrVq1+O53v1t5E/sWefzxxyOTycTjjz9e1VPJs3Tp0th+++2jpKTkG13vhRdeGJ06ddqs+3zb32dl75H77ruvqqdSaf785z9HJpOJ2bNnb3TZ1O8hAPi2EhoBYCt00003RSaT2ey4wYZNmzYtLrjggjjggANi3Lhxcfnll1f1lKrUTTfdFH/+85+rehqb7Lrrrot69erFgAEDvtH1DhkyJF599dWYMmXKJi3vfQYA/LeqXtUTAADKKykpidatW8fzzz8f7777brRr166qp7TVePvtt6OgYMv+X+ljjz0WBQUF8ac//Slq1qxZyTPb9tx0003RuHHjcntmde3aNVasWLFVbaPVq1fHddddF0OHDo1q1ap9o+tu1qxZHHvssXH11VdH7969N7q899m26ZRTTokBAwZEYWFhVU8FALZZ9mgEgK3MrFmz4plnnolrrrkmmjRp8o0fJhoRUVpaGitXrvzG17spCgsLo0aNGlt03/nz50etWrUqNf4sX7680h5ra1FQUBBFRUVbHHS/Dg8++GAsWLAg+vXrVyXr79evXzz11FPxn//8Z6PLVvb7LJvNxooVKyrlsb4Nvq7vp2rVqkVRUVFkMplKf+wN+fLLL+OLL774RtcJAF+Xree/HgGAiPhqb8aGDRtGz5494/jjj88LjatXr47tttsuTjvttHL3W7JkSRQVFcX555+fG1u1alUMHz482rVrF4WFhdGyZcu44IILYtWqVXn3zWQycdZZZ0VJSUnsscceUVhYGA8//HBERFx99dXRpUuXaNSoUdSqVSs6duxY4XnZVqxYEeecc040btw46tWrF7179465c+dGJpOJSy65JG/ZuXPnxumnnx5NmzaNwsLC2GOPPeJ//ud/Nmn7rHtutLLzqj399NNx3nnnRZMmTaJOnTrxgx/8IBYsWJD3HMeNGxfLli2LTCYTmUwm77Dhu+66Kzp27Bi1atWK7bbbLgYMGBBz5szJW3f37t1jzz33jJdeeim6du0atWvXjl//+tdbtK0nT54ce+65Z+75l23vdbfTGWecEd/5zneisLAw2rRpEz/72c/yosSiRYtiyJAh0bJlyygsLIx27drFlVdeGaWlpRvdjm+88UY88cQTue3RvXv3iKj4HI1lz/21116Lbt26Re3ataNdu3a598ITTzwRnTp1ilq1asWuu+4ajz76aIXPZ0tf98mTJ0fr1q2jbdu2eeOnnnpq1K1bNz744IPo1atX1K1bN3bYYYe48cYbIyLi9ddfj0MOOSTq1KkTrVq1irvvvjvv/qtXr44RI0bEzjvvHEVFRdGoUaM48MAD45FHHslbrkePHhER8de//nWD89zQ++zLL7+MSy+9NNq2bRuFhYXRunXr+PWvf13uPdK6devo1atXTJ06Nfbbb7+oVatW3HLLLRtc73PPPRdHHnlk1K9fP2rXrh3dunWLp59+Om+Z999/P84888zYddddo1atWtGoUaM44YQTKjwn4aJFi2Lo0KHRunXrKCwsjBYtWsSPfvSj+OSTT/KWKy0tjcsuuyxatGgRRUVFceihh8a77767wblGRFxyySWRyWTirbfein79+kVxcXE0atQozj333HIRcUPfT//85z/jqKOOiuLi4qhbt24ceuih8eyzz+bu++KLL0Ymk4nbb7+93BymTp0amUwmHnzwwYio+ByN2Ww2Ro4cGS1atIjatWvHwQcfHG+88UaFz2lTPouzZ8+OTCYTV199dYwZMyb3XnjzzTc3us0AYFvg0GkA2MqUlJTEcccdFzVr1owTTzwxbr755njhhRfie9/7XtSoUSN+8IMfxMSJE+OWW27J22Nq8uTJsWrVqtz560pLS6N3797x1FNPxeDBg6N9+/bx+uuvx7XXXhvvvPNOTJ48OW+9jz32WNxzzz1x1llnRePGjaN169YR8dV58Xr37h0nn3xyfPHFFzF+/Pg44YQT4sEHH4yePXvm7n/qqafGPffcE6ecckrsv//+8cQTT+TdXmbevHmx//775+JBkyZN4qGHHoozzjgjlixZEkOGDNmi7Xb22WdHw4YNY/jw4TF79uwYM2ZMnHXWWTFhwoSIiLjzzjvj1ltvjeeffz5uu+22iIjo0qVLRERcdtll8dvf/jb69esXgwYNigULFsQNN9wQXbt2jX/+85/RoEGD3HoWLlwYRx11VAwYMCB++MMfRtOmTTd7Wz/11FMxceLEOPPMM6NevXpx/fXXR9++feODDz6IRo0aRUTEhx9+GN///vdj0aJFMXjw4Nhtt91i7ty5cd9998Xy5cujZs2asXz58ujWrVvMnTs3fvKTn8SOO+4YzzzzTAwbNiw++uijGDNmzHq315gxY+Lss8+OunXrxkUXXRQREU2bNt3gNv7ss8+iV69eMWDAgDjhhBPi5ptvjgEDBkRJSUkMGTIkfvrTn8ZJJ50UV111VRx//PExZ86cqFevXkSkv+7PPPNM7LvvvhXetmbNmjjqqKOia9euMXr06CgpKYmzzjor6tSpExdddFGcfPLJcdxxx8XYsWPjRz/6UXTu3DnatGkTEV8Fr1GjRsWgQYPi+9//fixZsiRefPHFePnll+Owww7LraN+/frRtm3bePrpp2Po0KHrneeG3meDBg2K22+/PY4//vj4xS9+Ec8991yMGjUq/v3vf8ekSZPyHuftt9+OE088MX7yk5/Ej3/849h1113Xu87HHnssjjrqqOjYsWMMHz48CgoKYty4cXHIIYfEk08+Gd///vcjIuKFF16IZ555JgYMGBAtWrSI2bNnx8033xzdu3ePN998M2rXrh0RX11056CDDop///vfcfrpp8e+++4bn3zySUyZMiX+7//+Lxo3bpxb9xVXXBEFBQVx/vnnx+LFi2P06NFx8sknx3PPPbfe+a6tX79+0bp16xg1alQ8++yzcf3118dnn30Wd9xxR7nnuO730xtvvBEHHXRQFBcXxwUXXBA1atSIW265Jbp3754L3/vtt1/stNNOcc8998TAgQPzHnPChAnRsGHDOOKII9Y7v4svvjhGjhwZRx99dBx99NHx8ssvx+GHH15uD8TN/SyOGzcuVq5cGYMHD47CwsLYbrvtNml7AcBWLwsAbDVefPHFbERkH3nkkWw2m82WlpZmW7RokT333HNzy0ydOjUbEdkHHngg775HH310dqeddsr9feedd2YLCgqyTz75ZN5yY8eOzUZE9umnn86NRUS2oKAg+8Ybb5Sb0/Lly/P+/uKLL7J77rln9pBDDsmNvfTSS9mIyA4ZMiRv2VNPPTUbEdnhw4fnxs4444xs8+bNs5988knesgMGDMjWr1+/3PrW1apVq+zAgQNzf48bNy4bEdkePXpkS0tLc+NDhw7NVqtWLbto0aLc2MCBA7N16tTJe7zZs2dnq1Wrlr3sssvyxl9//fVs9erV88a7deuWjYjs2LFj85bd3G1ds2bN7Lvvvpsbe/XVV7MRkb3hhhtyYz/60Y+yBQUF2RdeeKHcNih7npdeemm2Tp062XfeeSfv9gsvvDBbrVq17AcffFDuvmvbY489st26dSs3Pn369GxEZKdPn17uud999925sbfeeiv33nn22Wdz42Xv0XHjxuXGUl731atXZzOZTPYXv/hFudsGDhyYjYjs5Zdfnhv77LPPsrVq1cpmMpns+PHjy8137ffjPvvsk+3Zs+d61722ww8/PNu+ffuNLlfR++yVV17JRkR20KBBeePnn39+NiKyjz32WG6sVatW2YjIPvzwwxtdV2lpaXbnnXfOHnHEEXnv/+XLl2fbtGmTPeyww/LG1jVjxoxsRGTvuOOO3NjFF1+cjYjsxIkTK1xfNvv/v0fat2+fXbVqVe726667LhsR2ddff32D8x4+fHg2IrK9e/fOGz/zzDOzEZF99dVXc2Pr+37q06dPtmbNmtn33nsvN/bhhx9m69Wrl+3atWtubNiwYdkaNWpkP/3009zYqlWrsg0aNMiefvrpubGy75JZs2Zls9lsdv78+dmaNWtme/bsmbdtf/3rX2cjIu97aFM/i7NmzcpGRLa4uDg7f/78DW4jANgWOXQaALYiJSUl0bRp0zj44IMj4qtDBvv37x/jx4+PNWvWRETEIYccEo0bN87tqRfx1Z5mjzzySPTv3z83du+990b79u1jt912i08++ST3c8ghh0RExPTp0/PW3a1bt9h9993LzalWrVp561m8eHEcdNBB8fLLL+fGyw5jPPPMM/Pue/bZZ+f9nc1m4/77749jjjkmstls3ryOOOKIWLx4cd7jbo7BgwfnnVvtoIMOijVr1sT777+/wftNnDgxSktLo1+/fnnzadasWey8887ltlNhYWG5Q9c3d1v36NEj7xDgvffeO4qLi3Pn/ystLY3JkyfHMcccE/vtt1+5OZc9z3vvvTcOOuigaNiwYd56e/ToEWvWrIl//OMfG9tsm6Vu3bp5V3zeddddo0GDBtG+ffu8K6SX/V72fFJf908//TSLine truncated
|
||||||
|
"text/plain": [
|
||||||
|
"<Figure size 1600x1200 with 1 Axes>"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"metadata": {},
|
||||||
|
"output_type": "display_data"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"%matplotlib inline\n",
|
||||||
|
"\n",
|
||||||
|
"import matplotlib\n",
|
||||||
|
"import matplotlib.pyplot as plt\n",
|
||||||
|
"import numpy as np\n",
|
||||||
|
"import os\n",
|
||||||
|
"\n",
|
||||||
|
"# Compute average inference time + std\n",
|
||||||
|
"time_results = {k: np.mean(v.model_inference_time) * 1e3 for k, v in results.items()}\n",
|
||||||
|
"time_results_std = np.std([v.model_inference_time for v in results.values()]) * 1000\n",
|
||||||
|
"\n",
|
||||||
|
"plt.rcdefaults()\n",
|
||||||
|
"fig, ax = plt.subplots(figsize=(16, 12))\n",
|
||||||
|
"ax.set_ylabel(\"Avg Inference time (ms)\")\n",
|
||||||
|
"ax.set_title(\"Average inference time (ms) for each provider\")\n",
|
||||||
|
"ax.bar(time_results.keys(), time_results.values(), yerr=time_results_std)\n",
|
||||||
|
"plt.show()"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"metadata": {
|
||||||
|
"accelerator": "GPU",
|
||||||
|
"colab": {
|
||||||
|
"collapsed_sections": [],
|
||||||
|
"name": "ONNX Overview",
|
||||||
|
"provenance": [],
|
||||||
|
"toc_visible": true
|
||||||
|
},
|
||||||
|
"kernelspec": {
|
||||||
|
"display_name": "Python 3",
|
||||||
|
"language": "python",
|
||||||
|
"name": "python3"
|
||||||
|
},
|
||||||
|
"language_info": {
|
||||||
|
"codemirror_mode": {
|
||||||
|
"name": "ipython",
|
||||||
|
"version": 3
|
||||||
|
},
|
||||||
|
"file_extension": ".py",
|
||||||
|
"mimetype": "text/x-python",
|
||||||
|
"name": "python",
|
||||||
|
"nbconvert_exporter": "python",
|
||||||
|
"pygments_lexer": "ipython3",
|
||||||
|
"version": "3.6.9"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"nbformat": 4,
|
||||||
|
"nbformat_minor": 1
|
||||||
|
}
|
||||||
+2
-1
@@ -10,9 +10,10 @@ Pull Request and we'll review it so it can be included here.
|
|||||||
## Hugging Face's notebooks :hugs:
|
## Hugging Face's notebooks :hugs:
|
||||||
|
|
||||||
| Notebook | Description | |
|
| Notebook | Description | |
|
||||||
|:----------|:-------------:|------:|
|
|:----------|:-------------|------:|
|
||||||
| [Getting Started Tokenizers](https://github.com/huggingface/transformers/blob/master/notebooks/01-training-tokenizers.ipynb) | How to train and use your very own tokenizer |[](https://colab.research.google.com/github/huggingface/transformers/blob/master/notebooks/01-training-tokenizers.ipynb) |
|
| [Getting Started Tokenizers](https://github.com/huggingface/transformers/blob/master/notebooks/01-training-tokenizers.ipynb) | How to train and use your very own tokenizer |[](https://colab.research.google.com/github/huggingface/transformers/blob/master/notebooks/01-training-tokenizers.ipynb) |
|
||||||
| [Getting Started Transformers](https://github.com/huggingface/transformers/blob/master/notebooks/02-transformers.ipynb) | How to easily start using transformers | [](https://colab.research.google.com/github/huggingface/transformers/blob/master/notebooks/02-transformers.ipynb) |
|
| [Getting Started Transformers](https://github.com/huggingface/transformers/blob/master/notebooks/02-transformers.ipynb) | How to easily start using transformers | [](https://colab.research.google.com/github/huggingface/transformers/blob/master/notebooks/02-transformers.ipynb) |
|
||||||
| [How to use Pipelines](https://github.com/huggingface/transformers/blob/master/notebooks/03-pipelines.ipynb) | Simple and efficient way to use State-of-the-Art models on downstream tasks through transformers | [](https://colab.research.google.com/github/huggingface/transformers/blob/master/notebooks/03-pipelines.ipynb) |
|
| [How to use Pipelines](https://github.com/huggingface/transformers/blob/master/notebooks/03-pipelines.ipynb) | Simple and efficient way to use State-of-the-Art models on downstream tasks through transformers | [](https://colab.research.google.com/github/huggingface/transformers/blob/master/notebooks/03-pipelines.ipynb) |
|
||||||
| [How to train a language model](https://github.com/huggingface/blog/blob/master/notebooks/01_how_to_train.ipynb)| Highlight all the steps to effectively train Transformer model on custom data | [](https://colab.research.google.com/github/huggingface/blog/blob/master/notebooks/01_how_to_train.ipynb)|
|
| [How to train a language model](https://github.com/huggingface/blog/blob/master/notebooks/01_how_to_train.ipynb)| Highlight all the steps to effectively train Transformer model on custom data | [](https://colab.research.google.com/github/huggingface/blog/blob/master/notebooks/01_how_to_train.ipynb)|
|
||||||
| [How to generate text](https://github.com/huggingface/blog/blob/master/notebooks/02_how_to_generate.ipynb)| How to use different decoding methods for language generation with transformers | [](https://colab.research.google.com/github/huggingface/blog/blob/master/notebooks/02_how_to_generate.ipynb)|
|
| [How to generate text](https://github.com/huggingface/blog/blob/master/notebooks/02_how_to_generate.ipynb)| How to use different decoding methods for language generation with transformers | [](https://colab.research.google.com/github/huggingface/blog/blob/master/notebooks/02_how_to_generate.ipynb)|
|
||||||
|
| [How to export model to ONNX](https://github.com/huggingface/blog/blob/master/notebooks/04-onnx-export.ipynb) | Highlight how to export and run inference workloads through ONNX |
|
||||||
@@ -67,8 +67,18 @@ extras = {}
|
|||||||
|
|
||||||
extras["mecab"] = ["mecab-python3"]
|
extras["mecab"] = ["mecab-python3"]
|
||||||
extras["sklearn"] = ["scikit-learn"]
|
extras["sklearn"] = ["scikit-learn"]
|
||||||
extras["tf"] = ["tensorflow"]
|
|
||||||
extras["tf-cpu"] = ["tensorflow-cpu"]
|
# keras2onnx and onnxconverter-common version is specific through a commit until 1.7.0 lands on pypi
|
||||||
|
extras["tf"] = [
|
||||||
|
"tensorflow",
|
||||||
|
"onnxconverter-common @ git+git://github.com/microsoft/onnxconverter-common.git@f64ca15989b6dc95a1f3507ff6e4c395ba12dff5#egg=onnxconverter-common",
|
||||||
|
"keras2onnx @ git+git://github.com/onnx/keras-onnx.git@cbdc75cb950b16db7f0a67be96a278f8d2953b48#egg=keras2onnx"
|
||||||
|
]
|
||||||
|
extras["tf-cpu"] = [
|
||||||
|
"tensorflow-cpu",
|
||||||
|
"onnxconverter-common @ git+git://github.com/microsoft/onnxconverter-common.git@f64ca15989b6dc95a1f3507ff6e4c395ba12dff5#egg=onnxconverter-common",
|
||||||
|
"keras2onnx @ git+git://github.com/onnx/keras-onnx.git@cbdc75cb950b16db7f0a67be96a278f8d2953b48#egg=keras2onnx"
|
||||||
|
]
|
||||||
extras["torch"] = ["torch"]
|
extras["torch"] = ["torch"]
|
||||||
|
|
||||||
extras["serving"] = ["pydantic", "uvicorn", "fastapi", "starlette"]
|
extras["serving"] = ["pydantic", "uvicorn", "fastapi", "starlette"]
|
||||||
|
|||||||
@@ -0,0 +1,212 @@
|
|||||||
|
from argparse import ArgumentParser
|
||||||
|
from itertools import takewhile
|
||||||
|
from os import listdir, makedirs
|
||||||
|
from os.path import abspath, dirname, exists
|
||||||
|
from typing import Dict, List, Optional, Tuple
|
||||||
|
|
||||||
|
from transformers import is_tf_available, is_torch_available
|
||||||
|
from transformers.pipelines import Pipeline, pipeline
|
||||||
|
from transformers.tokenization_utils import BatchEncoding
|
||||||
|
|
||||||
|
|
||||||
|
class OnnxConverterArgumentParser(ArgumentParser):
|
||||||
|
"""
|
||||||
|
Wraps all the script arguments supported to export transformers models to ONNX IR
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
super(OnnxConverterArgumentParser, self).__init__("ONNX Converter")
|
||||||
|
|
||||||
|
self.add_argument("--model", type=str, required=True, help="Model's id or path (ex: bert-base-cased)")
|
||||||
|
self.add_argument("--tokenizer", type=str, help="Tokenizer's id or path (ex: bert-base-cased)")
|
||||||
|
self.add_argument("--framework", type=str, choices=["pt", "tf"], help="Framework for loading the model")
|
||||||
|
self.add_argument("--opset", type=int, default=11, help="ONNX opset to use")
|
||||||
|
self.add_argument("--check-loading", action="store_true", help="Check ONNX is able to load the model")
|
||||||
|
self.add_argument("output")
|
||||||
|
|
||||||
|
|
||||||
|
def ensure_valid_input(model, tokens, input_names):
|
||||||
|
"""
|
||||||
|
Ensure input are presented in the correct order, without any None
|
||||||
|
Args:
|
||||||
|
model: The model used to forward the input data
|
||||||
|
tokens: BatchEncoding holding the input data
|
||||||
|
input_names: The name of the inputs
|
||||||
|
|
||||||
|
Returns: Tuple
|
||||||
|
|
||||||
|
"""
|
||||||
|
model_args_name = model.forward.__code__.co_varnames
|
||||||
|
model_args_pos = [(model_args_name.index(name) - 1, name) for name in input_names]
|
||||||
|
model_args = [None] * (max(map(lambda x: x[0], model_args_pos)) + 1)
|
||||||
|
|
||||||
|
for arg_pos, arg_name in model_args_pos:
|
||||||
|
model_args[arg_pos] = tokens[arg_name]
|
||||||
|
|
||||||
|
model_args = tuple(model_args) # Need to be ordered
|
||||||
|
return tuple(takewhile(lambda arg: arg is not None, model_args))
|
||||||
|
|
||||||
|
|
||||||
|
def infer_shapes(nlp: Pipeline, framework: str) -> Tuple[List[str], List[str], Dict, BatchEncoding]:
|
||||||
|
def build_shape_dict(tensor, is_input: bool, seq_len: int):
|
||||||
|
if isinstance(tensor, (tuple, list)):
|
||||||
|
return [build_shape_dict(t, is_input, seq_len) for t in tensor]
|
||||||
|
|
||||||
|
else:
|
||||||
|
# Let's assume batch is the first axis with only 1 element (~~ might not be always true ...)
|
||||||
|
axes = {[axis for axis, numel in enumerate(tensor.shape) if numel == 1][0]: "batch"}
|
||||||
|
if is_input:
|
||||||
|
if len(tensor.shape) == 2:
|
||||||
|
axes[1] = "sequence"
|
||||||
|
else:
|
||||||
|
raise ValueError("Unable to infer tensor axes ({})".format(len(tensor.shape)))
|
||||||
|
else:
|
||||||
|
seq_axes = [dim for dim, shape in enumerate(tensor.shape) if shape == seq_len]
|
||||||
|
axes.update({dim: "sequence" for dim in seq_axes})
|
||||||
|
|
||||||
|
return axes
|
||||||
|
|
||||||
|
tokens = nlp.tokenizer.encode_plus("This is a sample output", return_tensors=framework)
|
||||||
|
seq_len = tokens.input_ids.shape[-1]
|
||||||
|
outputs = nlp.model(**tokens) if framework == "pt" else nlp.model(tokens)
|
||||||
|
|
||||||
|
if not isinstance(outputs, (list, tuple)):
|
||||||
|
outputs = (outputs,)
|
||||||
|
|
||||||
|
# Generate input names & axes
|
||||||
|
input_vars = list(tokens.keys())
|
||||||
|
input_dynamic_axes = {k: build_shape_dict(v, True, seq_len) for k, v in tokens.items()}
|
||||||
|
|
||||||
|
# flatten potentially grouped outputs (past for gpt2, attentions)
|
||||||
|
outputs_flat = []
|
||||||
|
for output in outputs:
|
||||||
|
if isinstance(output, (tuple, list)):
|
||||||
|
outputs_flat.extend(output)
|
||||||
|
else:
|
||||||
|
outputs_flat.append(output)
|
||||||
|
|
||||||
|
# Generate output names & axes
|
||||||
|
output_names = ["output_{}".format(i) for i in range(len(outputs_flat))]
|
||||||
|
output_dynamic_axes = {k: build_shape_dict(v, False, seq_len) for k, v in zip(output_names, outputs_flat)}
|
||||||
|
|
||||||
|
# Create the aggregated axes representation
|
||||||
|
dynamic_axes = dict(input_dynamic_axes, **output_dynamic_axes)
|
||||||
|
return input_vars, output_names, dynamic_axes, tokens
|
||||||
|
|
||||||
|
|
||||||
|
def load_graph_from_args(framework: str, model: str, tokenizer: Optional[str] = None) -> Pipeline:
|
||||||
|
# If no tokenizer provided
|
||||||
|
if tokenizer is None:
|
||||||
|
tokenizer = model
|
||||||
|
|
||||||
|
print("Loading pipeline (model: {}, tokenizer: {})".format(model, tokenizer))
|
||||||
|
|
||||||
|
# Allocate tokenizer and model
|
||||||
|
return pipeline("feature-extraction", model=model, framework=framework)
|
||||||
|
|
||||||
|
|
||||||
|
def convert_pytorch(nlp: Pipeline, opset: int, output: str):
|
||||||
|
if not is_torch_available():
|
||||||
|
raise Exception("Cannot convert because PyTorch is not installed. Please install torch first.")
|
||||||
|
|
||||||
|
import torch
|
||||||
|
from torch.onnx import export
|
||||||
|
|
||||||
|
print("PyTorch: {}".format(torch.__version__))
|
||||||
|
|
||||||
|
with torch.no_grad():
|
||||||
|
input_names, output_names, dynamic_axes, tokens = infer_shapes(nlp, "pt")
|
||||||
|
model_args = ensure_valid_input(nlp.model, tokens, input_names)
|
||||||
|
|
||||||
|
export(
|
||||||
|
nlp.model,
|
||||||
|
model_args,
|
||||||
|
f=output,
|
||||||
|
input_names=input_names,
|
||||||
|
output_names=output_names,
|
||||||
|
dynamic_axes=dynamic_axes,
|
||||||
|
do_constant_folding=True,
|
||||||
|
use_external_data_format=True,
|
||||||
|
enable_onnx_checker=True,
|
||||||
|
opset_version=opset,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def convert_tensorflow(nlp: Pipeline, opset: int, output: str):
|
||||||
|
if not is_tf_available():
|
||||||
|
raise Exception(
|
||||||
|
"Cannot convert {} because TF is not installed. Please install torch first.".format(args.model)
|
||||||
|
)
|
||||||
|
|
||||||
|
print("/!\\ Please note TensorFlow doesn't support exporting model > 2Gb /!\\")
|
||||||
|
|
||||||
|
try:
|
||||||
|
import tensorflow as tf
|
||||||
|
from keras2onnx import convert_keras, save_model, __version__ as k2ov
|
||||||
|
|
||||||
|
print("TensorFlow: {}, keras2onnx: {}".format(tf.version.VERSION, k2ov))
|
||||||
|
|
||||||
|
# Build
|
||||||
|
input_names, output_names, dynamic_axes, tokens = infer_shapes(nlp, "tf")
|
||||||
|
|
||||||
|
# Forward
|
||||||
|
nlp.model.predict(tokens.data)
|
||||||
|
onnx_model = convert_keras(nlp.model, nlp.model.name, target_opset=opset)
|
||||||
|
save_model(onnx_model, output)
|
||||||
|
|
||||||
|
except ImportError as e:
|
||||||
|
raise Exception(
|
||||||
|
"Cannot import {} required to convert TF model to ONNX. Please install {} first.".format(e.name, e.name)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def convert(framework: str, model: str, output: str, opset: int, tokenizer: Optional[str] = None):
|
||||||
|
print("ONNX opset version set to: {}".format(opset))
|
||||||
|
|
||||||
|
# Load the pipeline
|
||||||
|
nlp = load_graph_from_args(framework, model, tokenizer)
|
||||||
|
|
||||||
|
parent = dirname(output)
|
||||||
|
if not exists(parent):
|
||||||
|
print("Creating folder {}".format(parent))
|
||||||
|
makedirs(parent)
|
||||||
|
elif len(listdir(parent)) > 0:
|
||||||
|
raise Exception("Folder {} is not empty, aborting conversion".format(parent))
|
||||||
|
|
||||||
|
# Export the graph
|
||||||
|
if framework == "pt":
|
||||||
|
convert_pytorch(nlp, opset, output)
|
||||||
|
else:
|
||||||
|
convert_tensorflow(nlp, opset, output)
|
||||||
|
|
||||||
|
|
||||||
|
def verify(path: str):
|
||||||
|
from onnxruntime import InferenceSession, SessionOptions
|
||||||
|
from onnxruntime.capi.onnxruntime_pybind11_state import RuntimeException
|
||||||
|
|
||||||
|
print("Checking ONNX model loading from: {}".format(path))
|
||||||
|
try:
|
||||||
|
onnx_options = SessionOptions()
|
||||||
|
_ = InferenceSession(path, onnx_options, providers=["CPUExecutionProvider"])
|
||||||
|
print("Model correctly loaded")
|
||||||
|
except RuntimeException as re:
|
||||||
|
print("Error while loading the model: {}".format(re))
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
parser = OnnxConverterArgumentParser()
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
# Make sure output is absolute path
|
||||||
|
args.output = abspath(args.output)
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Convert
|
||||||
|
convert(args.framework, args.model, args.output, args.opset, args.tokenizer)
|
||||||
|
|
||||||
|
# And verify
|
||||||
|
if args.check_loading:
|
||||||
|
verify(args.output)
|
||||||
|
except Exception as e:
|
||||||
|
print("Error while converting the model: {}".format(e))
|
||||||
|
exit(1)
|
||||||
@@ -0,0 +1,116 @@
|
|||||||
|
import unittest
|
||||||
|
from os import sep
|
||||||
|
from os.path import dirname, exists
|
||||||
|
from shutil import rmtree
|
||||||
|
|
||||||
|
from tests.utils import require_tf, require_torch
|
||||||
|
from transformers import BertConfig, BertTokenizerFast, FeatureExtractionPipeline
|
||||||
|
from transformers.convert_graph_to_onnx import convert, ensure_valid_input, infer_shapes
|
||||||
|
|
||||||
|
|
||||||
|
class FuncContiguousArgs:
|
||||||
|
def forward(self, input_ids, token_type_ids, attention_mask):
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
class FuncNonContiguousArgs:
|
||||||
|
def forward(self, input_ids, some_other_args, token_type_ids, attention_mask):
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
class OnnxExportTestCase(unittest.TestCase):
|
||||||
|
MODEL_TO_TEST = ["bert-base-cased", "gpt2", "roberta-base"]
|
||||||
|
|
||||||
|
@require_tf
|
||||||
|
def test_export_tensorflow(self):
|
||||||
|
for model in OnnxExportTestCase.MODEL_TO_TEST:
|
||||||
|
self._test_export(model, "tf", 11)
|
||||||
|
|
||||||
|
@require_torch
|
||||||
|
def test_export_pytorch(self):
|
||||||
|
for model in OnnxExportTestCase.MODEL_TO_TEST:
|
||||||
|
self._test_export(model, "pt", 11)
|
||||||
|
|
||||||
|
def _test_export(self, model, framework, opset):
|
||||||
|
try:
|
||||||
|
# Compute path
|
||||||
|
path = "onnx" + sep + model + ".onnx"
|
||||||
|
|
||||||
|
# Remove folder if exists
|
||||||
|
if exists(dirname(path)):
|
||||||
|
rmtree(dirname(path))
|
||||||
|
|
||||||
|
# Export
|
||||||
|
convert(framework, model, path, opset)
|
||||||
|
except Exception as e:
|
||||||
|
self.fail(e)
|
||||||
|
|
||||||
|
@require_torch
|
||||||
|
def test_infer_dynamic_axis_pytorch(self):
|
||||||
|
"""
|
||||||
|
Validate the dynamic axis generated for each parameters are correct
|
||||||
|
"""
|
||||||
|
from transformers import BertModel
|
||||||
|
|
||||||
|
model = BertModel(BertConfig.from_pretrained("bert-base-cased"))
|
||||||
|
tokenizer = BertTokenizerFast.from_pretrained("bert-base-cased")
|
||||||
|
self._test_infer_dynamic_axis(model, tokenizer, "pt")
|
||||||
|
|
||||||
|
@require_tf
|
||||||
|
def test_infer_dynamic_axis_tf(self):
|
||||||
|
"""
|
||||||
|
Validate the dynamic axis generated for each parameters are correct
|
||||||
|
"""
|
||||||
|
from transformers import TFBertModel
|
||||||
|
|
||||||
|
model = TFBertModel(BertConfig.from_pretrained("bert-base-cased"))
|
||||||
|
tokenizer = BertTokenizerFast.from_pretrained("bert-base-cased")
|
||||||
|
self._test_infer_dynamic_axis(model, tokenizer, "tf")
|
||||||
|
|
||||||
|
def _test_infer_dynamic_axis(self, model, tokenizer, framework):
|
||||||
|
nlp = FeatureExtractionPipeline(model, tokenizer)
|
||||||
|
|
||||||
|
variable_names = ["input_ids", "token_type_ids", "attention_mask", "output_0", "output_1"]
|
||||||
|
input_vars, output_vars, shapes, tokens = infer_shapes(nlp, framework)
|
||||||
|
|
||||||
|
# Assert all variables are present
|
||||||
|
self.assertEqual(len(shapes), len(variable_names))
|
||||||
|
self.assertTrue(all([var_name in shapes for var_name in variable_names]))
|
||||||
|
self.assertSequenceEqual(variable_names[:3], input_vars)
|
||||||
|
self.assertSequenceEqual(variable_names[3:], output_vars)
|
||||||
|
|
||||||
|
# Assert inputs are {0: batch, 1: sequence}
|
||||||
|
for var_name in ["input_ids", "token_type_ids", "attention_mask"]:
|
||||||
|
self.assertDictEqual(shapes[var_name], {0: "batch", 1: "sequence"})
|
||||||
|
|
||||||
|
# Assert outputs are {0: batch, 1: sequence} and {0: batch}
|
||||||
|
self.assertDictEqual(shapes["output_0"], {0: "batch", 1: "sequence"})
|
||||||
|
self.assertDictEqual(shapes["output_1"], {0: "batch"})
|
||||||
|
|
||||||
|
def test_ensure_valid_input(self):
|
||||||
|
"""
|
||||||
|
Validate parameters are correctly exported
|
||||||
|
GPT2 has "past" parameter in the middle of input_ids, token_type_ids and attention_mask.
|
||||||
|
ONNX doesn't support export with a dictionary, only a tuple. Thus we need to ensure we remove
|
||||||
|
token_type_ids and attention_mask for now to not having a None tensor in the middle
|
||||||
|
"""
|
||||||
|
# All generated args are valid
|
||||||
|
input_names = ["input_ids", "attention_mask", "token_type_ids"]
|
||||||
|
tokens = {"input_ids": [1, 2, 3, 4], "attention_mask": [0, 0, 0, 0], "token_type_ids": [1, 1, 1, 1]}
|
||||||
|
inputs_args = ensure_valid_input(FuncContiguousArgs(), tokens, input_names)
|
||||||
|
|
||||||
|
# Should have exactly the same number of args (all are valid)
|
||||||
|
self.assertEqual(len(inputs_args), 3)
|
||||||
|
|
||||||
|
# Parameter should be reordered according to their respective place in the function:
|
||||||
|
# (input_ids, token_type_ids, attention_mask)
|
||||||
|
self.assertEqual(inputs_args, (tokens["input_ids"], tokens["token_type_ids"], tokens["attention_mask"]))
|
||||||
|
|
||||||
|
# Generated args are interleaved with another args (for instance parameter "past" in GPT2)
|
||||||
|
inputs_args = ensure_valid_input(FuncNonContiguousArgs(), tokens, input_names)
|
||||||
|
|
||||||
|
# Should have exactly the one arg (all before the one not provided "some_other_args")
|
||||||
|
self.assertEqual(len(inputs_args), 1)
|
||||||
|
|
||||||
|
# Should have only "input_ids"
|
||||||
|
self.assertEqual(inputs_args[0], tokens["input_ids"])
|
||||||
Reference in new issue
Block a user