<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Accelerate BERT encoder with TF-TRT
Introduction
The NVIDIA TensorRT is a C++ library that facilitates high performance inference on NVIDIA graphics processing units (GPUs). TensorFlow™ integration with TensorRT™ (TF-TRT) optimizes TensorRT compatible parts of your computation graph, allowing TensorFlow to execute the remaining graph. While you can use TensorFlow's wide and flexible feature set, TensorRT will produce a highly optimized runtime engine for the TensorRT compatible subgraphs of your network.
In this notebook, we demonstrate accelerating BERT inference using TF-TRT. We focus on the encoder.
Requirements
This notebook requires at least TF 2.5 and TRT 7.1.3.
1. Download the model
We will download a bert base model from TF-Hub.
Step1:  2. Inference
In this section we will convert the model using TF-TRT and run inference.
Step2:  2.1 Helper functions
Step3:  2.2 Convert the model with TF-TRT
Step4:  2.3 Run inference with converted model
Step5:  Compare to the original function
Step6:  3. Dynamic sequence length
The sequence length for the encoder is dynamic, we can use different input sequence lengths. Here we call the original model for two sequences.
Step7:  The converted model is optimized for a sequnce length of 128 (and batch size 8). If we infer the converted model using a different sequence length, then two things can happen
<ASSISTANT_TASK:>
Python Code:
!pip install -q tf-models-official
import tensorflow as tf
import tensorflow_hub as hub
tfhub_handle_encoder = 'https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/3'
bert_saved_model_path = 'bert_base'
bert_model = hub.load(tfhub_handle_encoder)
tf.saved_model.save(bert_model, bert_saved_model_path)
Explanation: Accelerate BERT encoder with TF-TRT
Introduction
The NVIDIA TensorRT is a C++ library that facilitates high performance inference on NVIDIA graphics processing units (GPUs). TensorFlow™ integration with TensorRT™ (TF-TRT) optimizes TensorRT compatible parts of your computation graph, allowing TensorFlow to execute the remaining graph. While you can use TensorFlow's wide and flexible feature set, TensorRT will produce a highly optimized runtime engine for the TensorRT compatible subgraphs of your network.
In this notebook, we demonstrate accelerating BERT inference using TF-TRT. We focus on the encoder.
Requirements
This notebook requires at least TF 2.5 and TRT 7.1.3.
1. Download the model
We will download a bert base model from TF-Hub.
End of explanation
import matplotlib.pyplot as plt
import numpy as np
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import tag_constants
from tensorflow.python.compiler.tensorrt import trt_convert as trt
from timeit import default_timer as timer
tf.get_logger().setLevel('ERROR')
Explanation: 2. Inference
In this section we will convert the model using TF-TRT and run inference.
End of explanation
def get_func_from_saved_model(saved_model_dir):
    saved_model_loaded = tf.saved_model.load(
        saved_model_dir, tags=[tag_constants.SERVING])
    graph_func = saved_model_loaded.signatures[
        signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
    return graph_func, saved_model_loaded
def predict_and_benchmark_throughput(input_dict, model, N_warmup_run=50, N_run=500,
                                     result_key='predictions', batch_size=None):
    elapsed_time = []
    
    for val in input_dict.values():
        input_batch_size = val.shape[0]
        break
    if batch_size is None or batch_size > input_batch_size:
        batch_size = input_batch_size
        
    print('Benchmarking with batch size', batch_size)
    
    elapsed_time = np.zeros(N_run)
    for i in range(N_warmup_run):                                             
        preds = model(**input_dict)
    
    # Force device synchronization with .numpy()
    tmp = preds[result_key][0].numpy() 
    
    for i in range(N_run):
        start_time = timer()
        preds = model(**input_dict)
        # Synchronize
        tmp += preds[result_key][0].numpy() 
        end_time = timer()
        elapsed_time[i] = end_time - start_time
        if i>=50 and i % 50 == 0:
            print('Steps {}-{} average: {:4.1f}ms'.format(i-50, i, (elapsed_time[i-50:i].mean()) * 1000))
    latency = elapsed_time.mean() * 1000
    print('Latency: {:5.2f}+/-{:4.2f}ms'.format(latency, elapsed_time.std() * 1000))
    print('Throughput: {:.0f} samples/s'.format(N_run * batch_size / elapsed_time.sum()))
    return latency
def trt_convert(input_path, output_path, input_shapes, explicit_batch=False,
                dtype=np.float32, precision='FP32', prof_strategy='Optimal'):
    conv_params=trt.TrtConversionParams(
        precision_mode=precision, minimum_segment_size=50,
        max_workspace_size_bytes=12*1<<30, maximum_cached_engines=1)
    converter = trt.TrtGraphConverterV2(
        input_saved_model_dir=input_path, conversion_params=conv_params,
        use_dynamic_shape=explicit_batch,
        dynamic_shape_profile_strategy=prof_strategy)
    converter.convert()
    def input_fn():
        for shapes in input_shapes:
            # return a list of input tensors
            yield [np.ones(shape=x).astype(dtype) for x in shapes]
    converter.build(input_fn)
    converter.save(output_path)
    
def random_input(batch_size, seq_length):
    # Generate random input data
    mask = tf.convert_to_tensor(np.ones((batch_size, seq_length), dtype=np.int32))
    type_id = tf.convert_to_tensor(np.zeros((batch_size, seq_length), dtype=np.int32))
    word_id = tf.convert_to_tensor(np.random.randint(0, 1000, size=[batch_size, seq_length], dtype=np.int32))
    return {'input_mask':mask, 'input_type_ids': type_id, 'input_word_ids':word_id}
Explanation: 2.1 Helper functions
End of explanation
bert_trt_path = bert_saved_model_path + '_trt'
input_shapes = [[(1, 128), (1, 128), (1, 128)]] 
trt_convert(bert_saved_model_path, bert_trt_path, input_shapes, True, np.int32, precision='FP16')
Explanation: 2.2 Convert the model with TF-TRT
End of explanation
trt_func, _ = get_func_from_saved_model(bert_trt_path)
input_dict = random_input(1, 128)
result_key = 'bert_encoder_1' # 'classifier'
res = predict_and_benchmark_throughput(input_dict, trt_func, result_key=result_key)
Explanation: 2.3 Run inference with converted model
End of explanation
func, model = get_func_from_saved_model(bert_saved_model_path)
res = predict_and_benchmark_throughput(input_dict, func, result_key=result_key)
Explanation: Compare to the original function
End of explanation
seq1 = random_input(1, 128)
res1 = func(**seq1)
seq2 = random_input(1, 180)
res2 = func(**seq2)
Explanation: 3. Dynamic sequence length
The sequence length for the encoder is dynamic, we can use different input sequence lengths. Here we call the original model for two sequences.
End of explanation
bert_trt_path = bert_saved_model_path + '_trt2'
input_shapes = [[(1, 128), (1, 128), (1, 128)], [(1, 180), (1, 180), (1, 180)]] 
trt_convert(bert_saved_model_path, bert_trt_path, input_shapes, True, np.int32, precision='FP16',
            prof_strategy='Range')
trt_func_dynamic, _ = get_func_from_saved_model(bert_trt_path)
trt_res = trt_func_dynamic(**seq1)
result_key = 'bert_encoder_1' # 'classifier'
res = predict_and_benchmark_throughput(seq1, trt_func_dynamic, result_key=result_key)
res = predict_and_benchmark_throughput(seq2, trt_func_dynamic, result_key=result_key)
Explanation: The converted model is optimized for a sequnce length of 128 (and batch size 8). If we infer the converted model using a different sequence length, then two things can happen:
1. If TrtConversionParams.allow_build_at_runtime == False: native TF model is inferred
2. if TrtConversionParams.allow_build_at_runtime == True a new TRT engine is created which is optimized for the new sequence length. 
The first option do not provide TRT accelaration while the second one creates a large overhead while the new engine is constructed. In the next section we convert the model to handle multiple sequence lengths.
3.1 TRT Conversion with dynamic sequence length
End of explanation
<END_TASK>