<jupyter_start><jupyter_text>Nvidia GPU INT-8 quantization on any transformers model (encoder based) For some context and explanations, please check our documentation here: [https://els-rd.github.io/transformer-deploy/quantization/quantization_intro/](https://els-rd.github.io/transformer-deploy/quantization/quantization_intro/).  Project setup Dependencies installation Your machine should have Nvidia CUDA 11.X, TensorRT 8.2.1 and cuBLAS installed. It's said to be tricky to install, in my experience, just follow Nvidia download page instructions **and nothing else**, it should work out of the box. Nvidia Docker image could be a good choice too.<jupyter_code>#! pip3 install git+ssh://git@github.com/ELS-RD/transformer-deploy
#! pip3 install datasets sklearn
#! pip3 install git+ssh://git@github.com/NVIDIA/TensorRT#egg=pytorch-quantization\&subdirectory=tools/pytorch-quantization/<jupyter_output>Looking in indexes: https://pypi.org/simple, https://pypi.ngc.nvidia.com
Collecting datasets
  Downloading datasets-1.18.4-py3-none-any.whl (312 kB)
[K     |████████████████████████████████| 312 kB 12.2 MB/s eta 0:00:01
[?25hCollecting sklearn
  Downloading sklearn-0.0.tar.gz (1.1 kB)
Collecting dill
  Downloading dill-0.3.4-py2.py3-none-any.whl (86 kB)
[K     |████████████████████████████████| 86 kB 70.0 MB/s eta 0:00:01
[?25hRequirement already satisfied: requests>=2.19.0 in /home/geantvert/.local/share/virtualenvs/fast_transformer/lib/python3.9/site-packages (from datasets) (2.27.1)
Requirement already satisfied: numpy>=1.17 in /home/geantvert/.local/share/virtualenvs/fast_transformer/lib/python3.9/site-packages (from datasets) (1.21.5)
Requirement already satisfied: packaging in /home/geantvert/.local/share/virtualenvs/fast_transformer/lib/python3.9/site-packages (from datasets) (21.3)
Collecting pyarrow!=4.0.0,>=3.0.0
  Downloading pyarrow-7.0.0-cp39-cp39-manylinu[...]<jupyter_text>Check the GPU is enabled and usable.<jupyter_code>! nvidia-smi
import logging
import os
from collections import OrderedDict
from typing import Dict, List
from typing import OrderedDict as OD
from typing import Union

import datasets
import numpy as np
import tensorrt as trt
import torch
import transformers
from datasets import load_dataset, load_metric
from tensorrt.tensorrt import IExecutionContext, Logger, Runtime

from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    IntervalStrategy,
    PreTrainedModel,
    PreTrainedTokenizer,
    Trainer,
    TrainingArguments,
)

from transformer_deploy.backends.ort_utils import (
    cpu_quantization,
    create_model_for_provider,
    optimize_onnx,
)
from transformer_deploy.backends.pytorch_utils import convert_to_onnx
from transformer_deploy.backends.trt_utils import build_engine, get_binding_idxs, infer_tensorrt
from transformer_deploy.benchmarks.utils import print_timings, track_infer_time
from transformer_deploy.QDQModels.calibration_utils import QATCalibrate<jupyter_output><empty_output><jupyter_text>Set logging to `error` level to ease readability of this `notebook` on Github.<jupyter_code>log_level = logging.ERROR
logging.getLogger().setLevel(log_level)
datasets.utils.logging.set_verbosity(log_level)
transformers.utils.logging.set_verbosity(log_level)
transformers.utils.logging.enable_default_handler()
transformers.utils.logging.enable_explicit_format()
trt_logger: Logger = trt.Logger(trt.Logger.ERROR)
transformers.logging.set_verbosity_error()<jupyter_output><empty_output><jupyter_text>Preprocess data This part is inspired from an [official Notebooks from Hugging Face](https://github.com/huggingface/notebooks/blob/master/examples/text_classification.ipynb).There is nothing special to do. Define the task:<jupyter_code>model_name = "roberta-base"
task = "mnli"
num_labels = 3
batch_size = 32
max_seq_len = 256
validation_key = "validation_matched"
timings: Dict[str, List[float]] = dict()
runtime: Runtime = trt.Runtime(trt_logger)
profile_index = 0<jupyter_output><empty_output><jupyter_text>Preprocess data (task specific):<jupyter_code>def preprocess_function(examples):
    return tokenizer(
        examples["premise"], examples["hypothesis"], truncation=True, padding="max_length", max_length=max_seq_len
    )


def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    if task != "stsb":
        predictions = np.argmax(predictions, axis=1)
    else:
        predictions = predictions[:, 0]
    return metric.compute(predictions=predictions, references=labels)


def convert_tensor(data: OD[str, List[List[int]]], output: str) -> OD[str, Union[np.ndarray, torch.Tensor]]:
    input: OD[str, Union[np.ndarray, torch.Tensor]] = OrderedDict()
    for k in ["input_ids", "attention_mask", "token_type_ids"]:
        if k in data:
            v = data[k]
            if output == "torch":
                value = torch.tensor(v, dtype=torch.long, device="cuda")
            elif output == "np":
                value = np.asarray(v, dtype=np.int32)
            else:
                raise Exception(f"unknown output type: {output}")
            input[k] = value
    return input


def measure_accuracy(infer, tensor_type: str) -> float:
    outputs = list()
    for start_index in range(0, len(encoded_dataset[validation_key]), batch_size):
        end_index = start_index + batch_size
        data = encoded_dataset[validation_key][start_index:end_index]
        inputs: OD[str, np.ndarray] = convert_tensor(data=data, output=tensor_type)
        output = infer(inputs)[0]
        if tensor_type == "torch":
            output = output.detach().cpu().numpy()
        output = np.argmax(output, axis=1).astype(int).tolist()
        outputs.extend(output)
    return np.mean(np.array(outputs) == np.array(validation_labels))


def get_trainer(model: PreTrainedModel) -> Trainer:
    trainer = Trainer(
        model,
        args,
        train_dataset=encoded_dataset["train"],
        eval_dataset=encoded_dataset[validation_key],
        tokenizer=tokenizer,
        compute_metrics=compute_metrics,
    )
    transformers.logging.set_verbosity_error()
    return trainer
tokenizer: PreTrainedTokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)

dataset = load_dataset("glue", task)
metric = load_metric("glue", task)

encoded_dataset = dataset.map(preprocess_function, batched=True)
validation_labels = [item["label"] for item in encoded_dataset[validation_key]]

nb_step = 1000
strategy = IntervalStrategy.STEPS
args = TrainingArguments(
    f"{model_name}-{task}",
    evaluation_strategy=strategy,
    eval_steps=nb_step,
    logging_steps=nb_step,
    save_steps=nb_step,
    save_strategy=strategy,
    learning_rate=1e-5,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size * 2,
    num_train_epochs=1,
    fp16=True,
    group_by_length=True,
    weight_decay=0.01,
    load_best_model_at_end=True,
    metric_for_best_model="accuracy",
    report_to=[],
)<jupyter_output><empty_output><jupyter_text>(Standard) fine-tuning modelNow that our data are ready, we can download/fine tune the pretrained model.<jupyter_code>model_fp16: PreTrainedModel = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=num_labels)
trainer = get_trainer(model_fp16)
transformers.logging.set_verbosity_error()
trainer.train()
print(trainer.evaluate())

model_fp16.save_pretrained("model_trained_fp16")<jupyter_output>[INFO|trainer.py:457] 2022-03-09 19:15:29,964 >> Using amp half precision backend
/home/geantvert/.local/share/virtualenvs/fast_transformer/lib/python3.9/site-packages/transformers/optimization.py:306: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning
  warnings.warn(<jupyter_text>Add quantization support to any modelThe idea is to take the source code of a specific model and add automatically `QDQ` nodes. QDQ nodes will be placed before and after an operation that we want to quantize, that’s inside these nodes that the information to perform the mapping between high precision and low precision number is stored.If you want to know more, check our documentation on: [https://els-rd.github.io/transformer-deploy/quantization/quantization_ast/](https://els-rd.github.io/transformer-deploy/quantization/quantization_ast/)<jupyter_code>for percentile in [99.9, 99.99, 99.999, 99.9999]:
    with QATCalibrate(method="histogram", percentile=percentile) as qat:
        model_q: PreTrainedModel = AutoModelForSequenceClassification.from_pretrained(
            "model_trained_fp16", num_labels=num_labels
        )
        model_q = model_q.cuda()
        qat.setup_model_qat(model_q)  # prepare quantizer to any model

        with torch.no_grad():
            for start_index in range(0, 128, batch_size):
                end_index = start_index + batch_size
                data = encoded_dataset["train"][start_index:end_index]
                input_torch = {
                    k: torch.tensor(v, dtype=torch.long, device="cuda")
                    for k, v in data.items()
                    if k in ["input_ids", "attention_mask", "token_type_ids"]
                }
                model_q(**input_torch)
    trainer = get_trainer(model_q)
    print(f"percentile: {percentile}")
    print(trainer.evaluate())<jupyter_output>[INFO|trainer.py:457] 2022-03-09 20:04:30,756 >> Using amp half precision backend<jupyter_text>As you can see, the chosen percentile value has a high impact on the final accuracy.For the rest of the notebook, we apply the `99.999` percentile.<jupyter_code>with QATCalibrate(method="histogram", percentile=99.999) as qat:
    model_q: PreTrainedModel = AutoModelForSequenceClassification.from_pretrained(
        "model_trained_fp16", num_labels=num_labels
    )
    model_q = model_q.cuda()
    qat.setup_model_qat(model_q)  # prepare quantizer to any model

    with torch.no_grad():
        for start_index in range(0, 128, batch_size):
            end_index = start_index + batch_size
            data = encoded_dataset["train"][start_index:end_index]
            input_torch = {
                k: torch.tensor(v, dtype=torch.long, device="cuda")
                for k, v in data.items()
                if k in ["input_ids", "attention_mask", "token_type_ids"]
            }
            model_q(**input_torch)
trainer = get_trainer(model_q)
print(trainer.evaluate())<jupyter_output>[INFO|trainer.py:457] 2022-03-09 20:21:33,073 >> Using amp half precision backend<jupyter_text>Per layer quantization analysisBelow we will run a sensitivity analysis, by enabling quantization of one layer at a time and measuring the accuracy. That way we will be able to detect if the quantization of a specific layer has a larger cost on accuracy than other layers.<jupyter_code>from pytorch_quantization import nn as quant_nn

for i in range(12):
    layer_name = f"layer.{i}"
    print(layer_name)
    for name, module in model_q.named_modules():
        if isinstance(module, quant_nn.TensorQuantizer):
            if layer_name in name:
                module.enable_quant()
            else:
                module.disable_quant()
    trainer.evaluate()
    print("----")<jupyter_output>layer.0
{'eval_loss': 0.34956735372543335, 'eval_accuracy': 0.86571574121243, 'eval_runtime': 20.7064, 'eval_samples_per_second': 474.009, 'eval_steps_per_second': 7.437}
----
layer.1
{'eval_loss': 0.3523275852203369, 'eval_accuracy': 0.8649006622516556, 'eval_runtime': 25.4843, 'eval_samples_per_second': 385.14, 'eval_steps_per_second': 6.043}
----
layer.2
{'eval_loss': 0.356509268283844, 'eval_accuracy': 0.8622516556291391, 'eval_runtime': 20.7496, 'eval_samples_per_second': 473.021, 'eval_steps_per_second': 7.422}
----
layer.3
{'eval_loss': 0.36036217212677, 'eval_accuracy': 0.8617422312786551, 'eval_runtime': 20.7815, 'eval_samples_per_second': 472.296, 'eval_steps_per_second': 7.41}
----
layer.4
{'eval_loss': 0.35000357031822205, 'eval_accuracy': 0.8643912379011717, 'eval_runtime': 20.7921, 'eval_samples_per_second': 472.053, 'eval_steps_per_second': 7.407}
----
layer.5
{'eval_loss': 0.354992538690567, 'eval_accuracy': 0.8644931227712684, 'eval_runtime': 20.7938, 'eval_samples_per[...]<jupyter_text>It seems that quantization of layers 2 to 6 has the largest accuracy impact. Operator quantization analysisBelow we will run a sensitivity analysis, by enabling quantization of one operator type at a time and measuring the accuracy. That way we will be able to detect if a specific operator has a larger cost on accuracy. On Roberta we only quantize `matmul` and `LayerNorm`, so we test both candidates.<jupyter_code>for op in ["matmul", "layernorm"]:
    for name, module in model_q.named_modules():
        if isinstance(module, quant_nn.TensorQuantizer):
            if op in name:
                module.enable_quant()
            else:
                module.disable_quant()
    print(op)
    trainer.evaluate()
    print("----")<jupyter_output>matmul
{'eval_loss': 0.3494793176651001, 'eval_accuracy': 0.8654100866021396, 'eval_runtime': 26.892, 'eval_samples_per_second': 364.978, 'eval_steps_per_second': 5.727}
----
layernorm
{'eval_loss': 0.3585323095321655, 'eval_accuracy': 0.8587875700458482, 'eval_runtime': 24.0982, 'eval_samples_per_second': 407.293, 'eval_steps_per_second': 6.391}
----<jupyter_text>It appears that `LayerNorm` quantization has a significant accuracy cost.Our goal is to disable quantization for as few operations as possible while preserving accuracy as much as possible. Therefore we will try to only disable quantization for `LayerNorm` on Layers 2 to 6.<jupyter_code>disable_layer_names = ["layer.2", "layer.3", "layer.4", "layer.6"]

for name, module in model_q.named_modules():
    if isinstance(module, quant_nn.TensorQuantizer):
        if any([f"{l}.output.layernorm" in name for l in disable_layer_names]):
            print(f"disable {name}")
            module.disable_quant()
        else:
            module.enable_quant()
trainer.evaluate()<jupyter_output>disable roberta.encoder.layer.2.output.layernorm_quantizer_0
disable roberta.encoder.layer.2.output.layernorm_quantizer_1
disable roberta.encoder.layer.3.output.layernorm_quantizer_0
disable roberta.encoder.layer.3.output.layernorm_quantizer_1
disable roberta.encoder.layer.4.output.layernorm_quantizer_0
disable roberta.encoder.layer.4.output.layernorm_quantizer_1
disable roberta.encoder.layer.6.output.layernorm_quantizer_0
disable roberta.encoder.layer.6.output.layernorm_quantizer_1
{'eval_loss': 0.3617263436317444, 'eval_accuracy': 0.8614365766683647, 'eval_runtime': 46.0379, 'eval_samples_per_second': 213.194, 'eval_steps_per_second': 3.345}<jupyter_text>By just disabling quantization for a single operator on a few layers, we keep most of the performance boost (quantization) but retrieve more than 1 point of accuracy. It's also possible to perform an analysis per quantizer to get a smaller granularity but it's a bit slow to run.If we stop here, it's called a Post Training Quantization (PTQ). Below, we will try to retrieve even more accuracy.  Quantization Aware Training (QAT) We retrain the model with 1/10 or 1/100 of the original learning rate. Our goal is to retrieve most of the original accuracy.<jupyter_code>args.learning_rate = 1e-7
trainer = get_trainer(model_q)
trainer.train()
print(trainer.evaluate())
model_q.save_pretrained("model-qat")<jupyter_output>[INFO|trainer.py:457] 2022-03-09 20:28:11,049 >> Using amp half precision backend
/home/geantvert/.local/share/virtualenvs/fast_transformer/lib/python3.9/site-packages/transformers/optimization.py:306: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning
  warnings.warn(<jupyter_text>Export a `QDQ Pytorch` model to `ONNX`We need to enable fake quantization mode from Pytorch.<jupyter_code>data = encoded_dataset["train"][1:3]
input_torch = convert_tensor(data, output="torch")
convert_to_onnx(
    model_pytorch=model_q,
    output_path="model_qat.onnx",
    inputs_pytorch=input_torch,
    quantization=True,
    var_output_seq=False,
)
del model_q
QATCalibrate.restore()<jupyter_output><empty_output><jupyter_text>Benchmark  Convert `ONNX` graph to `TensorRT` engine<jupyter_code>engine = build_engine(
    runtime=runtime,
    onnx_file_path="model_qat.onnx",
    logger=trt_logger,
    min_shape=(1, max_seq_len),
    optimal_shape=(batch_size, max_seq_len),
    max_shape=(batch_size, max_seq_len),
    workspace_size=10000 * 1024 * 1024,
    fp16=True,
    int8=True,
)
# same as above, but from the terminal
# !/usr/src/tensorrt/bin/trtexec --onnx=model_qat.onnx --shapes=input_ids:32x256,attention_mask:32x256 --best --workspace=10000 --saveEngine="test.plan"<jupyter_output><empty_output><jupyter_text>Prepare input and output buffer<jupyter_code>context: IExecutionContext = engine.create_execution_context()
context.set_optimization_profile_async(
    profile_index=profile_index, stream_handle=torch.cuda.current_stream().cuda_stream
)
input_binding_idxs, output_binding_idxs = get_binding_idxs(engine, profile_index)  # type: List[int], List[int]
data = encoded_dataset["train"][0:batch_size]
input_torch: OD[str, torch.Tensor] = convert_tensor(data=data, output="torch")
input_np: OD[str, np.ndarray] = convert_tensor(data=data, output="np")<jupyter_output><empty_output><jupyter_text>Inference on `TensorRT`We first check that inference is working correctly:<jupyter_code>tensorrt_output = infer_tensorrt(
    context=context,
    host_inputs=input_torch,
    input_binding_idxs=input_binding_idxs,
    output_binding_idxs=output_binding_idxs,
)
print(tensorrt_output)<jupyter_output>[tensor([[ 1.2287,  1.3706, -2.4623],
        [ 2.4742, -0.7816, -1.8427],
        [ 2.4837, -0.3966, -2.2666],
        [ 2.9077, -0.3062, -2.9778],
        [ 2.3437,  0.0488, -2.6377],
        [ 3.7914, -1.1918, -3.1387],
        [-3.6134,  2.7432,  0.8490],
        [ 3.6679, -1.5787, -2.6408],
        [ 1.0155, -1.2787,  0.4250],
        [-3.4514, -0.4434,  4.2748],
        [ 3.5201, -1.1297, -2.8988],
        [-3.0225, -0.4062,  3.8606],
        [-2.7311,  3.5470, -0.4632],
        [-2.0741,  1.6613,  0.5798],
        [-0.4047, -0.8650,  1.6144],
        [ 2.8432, -1.3301, -1.8994],
        [ 3.7722, -0.9103, -3.3070],
        [-2.4204, -2.1432,  4.6537],
        [-3.1179, -1.3207,  4.6400],
        [-1.8794,  4.1075, -1.8630],
        [ 3.7726, -1.2056, -3.0701],
        [ 1.8645,  1.9744, -3.8743],
        [-3.1448, -1.2497,  4.5782],
        [ 3.5385, -0.2421, -3.6629],
        [ 3.7501, -1.6469, -2.7108],
        [-0.6568,  0.9046, -0.0228],
        [-3.2998,  0.0867,  3.3673],
[...]<jupyter_text>Measure of the accuracy:<jupyter_code>infer_trt = lambda inputs: infer_tensorrt(
    context=context,
    host_inputs=inputs,
    input_binding_idxs=input_binding_idxs,
    output_binding_idxs=output_binding_idxs,
)

measure_accuracy(infer=infer_trt, tensor_type="torch")<jupyter_output><empty_output><jupyter_text>Latency measures:<jupyter_code>time_buffer = list()
for _ in range(100):
    with track_infer_time(time_buffer):
        _ = infer_tensorrt(
            context=context,
            host_inputs=input_torch,
            input_binding_idxs=input_binding_idxs,
            output_binding_idxs=output_binding_idxs,
        )

print_timings(name="TensorRT (INT-8)", timings=time_buffer)
del engine, context<jupyter_output><empty_output><jupyter_text>Pytorch baselineTime to get some numbers to compare with. GPU executionWe will measure vanilla Pytorch inference on both FP32 and FP16 precision on GPU, it will be our baseline:<jupyter_code>baseline_model = AutoModelForSequenceClassification.from_pretrained("model_trained_fp16", num_labels=num_labels)
baseline_model = baseline_model.cuda()
baseline_model = baseline_model.eval()

data = encoded_dataset["train"][0:batch_size]
input_torch: OD[str, torch.Tensor] = convert_tensor(data=data, output="torch")

with torch.inference_mode():
    for _ in range(30):
        _ = baseline_model(**input_torch)
        torch.cuda.synchronize()
    time_buffer = list()
    for _ in range(100):
        with track_infer_time(time_buffer):
            _ = baseline_model(**input_torch)
            torch.cuda.synchronize()
print_timings(name="Pytorch (FP32)", timings=time_buffer)
with torch.inference_mode():
    with torch.cuda.amp.autocast():
        for _ in range(30):
            _ = baseline_model(**input_torch)
            torch.cuda.synchronize()
        time_buffer = []
        for _ in range(100):
            with track_infer_time(time_buffer):
                _ = baseline_model(**input_torch)
                torch.cuda.synchronize()
print_timings(name="Pytorch (FP16)", timings=time_buffer)
del baseline_model<jupyter_output>[Pytorch (FP16)] mean=58.73ms, sd=1.88ms, min=55.56ms, max=69.69ms, median=58.03ms, 95p=62.28ms, 99p=64.42ms<jupyter_text>CPU execution<jupyter_code>baseline_model = AutoModelForSequenceClassification.from_pretrained("model_trained_fp16", num_labels=num_labels)
baseline_model = baseline_model.eval()
data = encoded_dataset["train"][0:batch_size]
input_torch: OD[str, torch.Tensor] = convert_tensor(data=data, output="torch")
input_torch_cpu = {k: v.to("cpu") for k, v in input_torch.items()}

torch.set_num_threads(os.cpu_count())

with torch.inference_mode():
    for _ in range(3):
        _ = baseline_model(**input_torch_cpu)
        torch.cuda.synchronize()
    time_buffer = list()
    for _ in range(10):
        with track_infer_time(time_buffer):
            _ = baseline_model(**input_torch_cpu)
            torch.cuda.synchronize()
print_timings(name="Pytorch (FP32) - CPU", timings=time_buffer)
with torch.inference_mode():
    with torch.cuda.amp.autocast():
        for _ in range(3):
            _ = baseline_model(**input_torch_cpu)
            torch.cuda.synchronize()
        time_buffer = []
        for _ in range(10):
            with track_infer_time(time_buffer):
                _ = baseline_model(**input_torch_cpu)
                torch.cuda.synchronize()
print_timings(name="Pytorch (FP16) - CPU", timings=time_buffer)
del baseline_model<jupyter_output>[Pytorch (FP16) - CPU] mean=4422.39ms, sd=170.45ms, min=4250.10ms, max=4744.97ms, median=4337.85ms, 95p=4727.72ms, 99p=4741.52ms<jupyter_text>Below, we will perform dynamic quantization on CPU.<jupyter_code>quantized_baseline_model = AutoModelForSequenceClassification.from_pretrained(
    "model_trained_fp16", num_labels=num_labels
)
quantized_baseline_model = quantized_baseline_model.eval()
quantized_baseline_model = torch.quantization.quantize_dynamic(
    quantized_baseline_model, {torch.nn.Linear}, dtype=torch.qint8
)

with torch.inference_mode():
    for _ in range(3):
        _ = quantized_baseline_model(**input_torch_cpu)
        torch.cuda.synchronize()
    time_buffer = list()
    for _ in range(10):
        with track_infer_time(time_buffer):
            _ = quantized_baseline_model(**input_torch_cpu)
            torch.cuda.synchronize()
print_timings(name="Pytorch (INT-8) - CPU", timings=time_buffer)<jupyter_output>[Pytorch (INT-8) - CPU] mean=3818.99ms, sd=137.98ms, min=3616.11ms, max=4049.00ms, median=3807.33ms, 95p=4024.45ms, 99p=4044.09ms<jupyter_text>TensorRT baseline Below we export our finetuned model, the purpose is to only check the performance on mixed precision (FP16, no quantization).<jupyter_code>baseline_model = AutoModelForSequenceClassification.from_pretrained("model_trained_fp16", num_labels=num_labels)
baseline_model = baseline_model.cuda()
convert_to_onnx(
    baseline_model, output_path="baseline.onnx", inputs_pytorch=input_torch, quantization=False, var_output_seq=False
)
del baseline_model
engine = build_engine(
    runtime=runtime,
    onnx_file_path="baseline.onnx",
    logger=trt_logger,
    min_shape=(batch_size, max_seq_len),
    optimal_shape=(batch_size, max_seq_len),
    max_shape=(batch_size, max_seq_len),
    workspace_size=10000 * 1024 * 1024,
    fp16=True,
    int8=False,
)
input_torch: OD[str, np.ndarray] = convert_tensor(data=data, output="torch")
context: IExecutionContext = engine.create_execution_context()
context.set_optimization_profile_async(
    profile_index=profile_index, stream_handle=torch.cuda.current_stream().cuda_stream
)
input_binding_idxs, output_binding_idxs = get_binding_idxs(engine, profile_index)  # type: List[int], List[int]
for _ in range(30):
    _ = infer_tensorrt(
        context=context,
        host_inputs=input_torch,
        input_binding_idxs=input_binding_idxs,
        output_binding_idxs=output_binding_idxs,
    )
time_buffer = list()
for _ in range(100):
    with track_infer_time(time_buffer):
        _ = infer_tensorrt(
            context=context,
            host_inputs=input_torch,
            input_binding_idxs=input_binding_idxs,
            output_binding_idxs=output_binding_idxs,
        )

print_timings(name="TensorRT (FP16)", timings=time_buffer)
del engine, context<jupyter_output>[TensorRT (FP16)] mean=32.36ms, sd=1.47ms, min=29.92ms, max=38.08ms, median=32.47ms, 95p=34.38ms, 99p=36.50ms<jupyter_text>ONNX Runtime baselineONNX Runtime is the go to inference solution from Microsoft.The recent 1.10 version of ONNX Runtime (with TensorRT support) is still a bit buggy on transformer models, that is why we use the 1.9.0 version in the measures below.As before, CPU quantization is dynamic.Function `` will set ONNX Runtime to use all cores available and enable any possible optimizations.<jupyter_code>optimize_onnx(
    onnx_path="baseline.onnx",
    onnx_optim_model_path="baseline-optimized.onnx",
    fp16=True,
    use_cuda=True,
    num_attention_heads=12,
    hidden_size=768,
    architecture="bert",
)
cpu_quantization(input_model_path="baseline.onnx", output_model_path="baseline-quantized.onnx")
labels = [item["label"] for item in encoded_dataset[validation_key]]
data = encoded_dataset[validation_key][0:batch_size]
inputs_onnx: OD[str, np.ndarray] = convert_tensor(data=data, output="np")

model = create_model_for_provider(path="baseline-optimized.onnx", provider_to_use="CUDAExecutionProvider")
output = model.run(None, inputs_onnx)
data = encoded_dataset["train"][0:batch_size]
inputs_onnx: OD[str, np.ndarray] = convert_tensor(data=data, output="np")

for provider, model_path, benchmark_name, warmup, nb_inference in [
    ("CUDAExecutionProvider", "baseline.onnx", "ONNX Runtime GPU (FP32)", 10, 100),
    ("CUDAExecutionProvider", "baseline-optimized.onnx", "ONNX Runtime GPU (FP16)", 10, 100),
    ("CPUExecutionProvider", "baseline.onnx", "ONNX Runtime CPU (FP32)", 3, 10),
    ("CPUExecutionProvider", "baseline-optimized.onnx", "ONNX Runtime CPU (FP16)", 3, 10),
    ("CPUExecutionProvider", "baseline-quantized.onnx", "ONNX Runtime CPU (INT-8)", 3, 10),
]:
    model = create_model_for_provider(path=model_path, provider_to_use=provider)
    for _ in range(warmup):
        _ = model.run(None, inputs_onnx)
    time_buffer = []
    for _ in range(nb_inference):
        with track_infer_time(time_buffer):
            _ = model.run(None, inputs_onnx)
    print_timings(name=benchmark_name, timings=time_buffer)
    del model<jupyter_output>[ONNX Runtime GPU (FP32)] mean=82.50ms, sd=11.53ms, min=74.12ms, max=118.73ms, median=77.62ms, 95p=112.51ms, 99p=117.81ms
[ONNX Runtime GPU (FP16)] mean=36.23ms, sd=4.01ms, min=33.47ms, max=55.75ms, median=35.35ms, 95p=46.89ms, 99p=53.95ms
[ONNX Runtime CPU (FP32)] mean=4608.31ms, sd=468.16ms, min=3954.09ms, max=5226.20ms, median=4531.47ms, 95p=5189.71ms, 99p=5218.90ms
[ONNX Runtime CPU (FP16)] mean=3987.95ms, sd=223.00ms, min=3755.70ms, max=4548.57ms, median=3907.74ms, 95p=4375.23ms, 99p=4513.90ms
[ONNX Runtime CPU (INT-8)] mean=3506.39ms, sd=125.07ms, min=3425.41ms, max=3872.95ms, median=3463.65ms, 95p=3713.08ms, 99p=3840.98ms<jupyter_text>Measure of the accuracy with ONNX Runtime engine and CUDA provider:<jupyter_code>model = create_model_for_provider(path="baseline.onnx", provider_to_use="CUDAExecutionProvider")
infer_ort = lambda tokens: model.run(None, tokens)
measure_accuracy(infer=infer_ort, tensor_type="np")
model = create_model_for_provider(path="baseline-optimized.onnx", provider_to_use="CUDAExecutionProvider")
infer_ort = lambda tokens: model.run(None, tokens)
measure_accuracy(infer=infer_ort, tensor_type="np")
model = create_model_for_provider(path="baseline-quantized.onnx", provider_to_use="CPUExecutionProvider")
infer_ort = lambda tokens: model.run(None, tokens)
measure_accuracy(infer=infer_ort, tensor_type="np")
del model<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11540733020670488, 'fraction_numerical': 0.0574604828903943, 'mean_word_length': 4.57544063528956, 'pattern_counts': {'":': 4, '<': 83, '<?xml version=': 0, '>': 97, 'https://': 7, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 2, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '24', 'dir': 'jupyter-structured-clean-dedup', 'id': '634486', 'license': 'Apache-2.0', 'n_tokens_mistral': 10844, 'n_tokens_neox': 9754, 'n_words': 2008, 'path': 'demo/quantization/quantization_end_to_end.ipynb', 'repo_name': 'dumpmemory/transformer-deploy', 'size': '28785'}