Introduction

If you have trained a deep learning model in TensorFlow 2.0 and want to deploy it with maximum inference speed on NVIDIA GPUs, converting your SavedModel to TensorRT is one of the most effective optimizations available. TensorRT is NVIDIA's high-performance inference optimizer and runtime that can deliver significantly lower latency and higher throughput compared to running the model directly through TensorFlow. This guide walks you through the conversion process step by step, including code examples and practical considerations.

Why Convert to TensorRT?

TensorRT applies a suite of optimizations to your model graph, including layer fusion, kernel auto-tuning, precision calibration, and memory reuse. These optimizations are tailored to the specific GPU hardware you are deploying on, meaning the resulting engine is purpose-built for your target device. For production inference workloads, this can reduce latency by 2x to 5x compared to a standard TensorFlow SavedModel.

Step 1 -- Save Your TensorFlow Model

Before converting, you need a model saved in the TensorFlow SavedModel format. Here is a simple example that creates and saves a model:

1import tensorflow as tf
2
3# Build a simple model
4model = tf.keras.Sequential([
5    tf.keras.layers.Dense(128, activation='relu', input_shape=(784,)),
6    tf.keras.layers.Dense(64, activation='relu'),
7    tf.keras.layers.Dense(10, activation='softmax')
8])
9
10model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
11
12# Save the model in SavedModel format
13model.save('my_saved_model')

The directory my_saved_model will contain the saved_model.pb file along with a variables/ directory. This is the input format that the TensorRT converter expects.

Step 2 -- Convert Using TF-TRT

TensorFlow ships with a built-in TensorRT integration called TF-TRT (TensorFlow-TensorRT). This is the recommended path for converting SavedModel to TensorRT. The TrtGraphConverterV2 class handles the conversion:

1from tensorflow.python.compiler.tensorrt import trt_convert as trt
2
3# Configure the converter
4converter = trt.TrtGraphConverterV2(
5    input_saved_model_dir='my_saved_model',
6    precision_mode=trt.TrtPrecisionMode.FP16,
7    max_workspace_size_bytes=1 << 30  # 1 GB
8)
9
10# Perform the conversion
11converter.convert()
12
13# Save the TensorRT-optimized model
14converter.save('my_trt_model')

The precision_mode parameter accepts FP32, FP16, or INT8. Using FP16 typically provides a good balance between speed and accuracy on modern NVIDIA GPUs. The max_workspace_size_bytes parameter controls how much GPU memory TensorRT can use during optimization.

Step 3 -- Run Inference with the Converted Model

Once the model is converted and saved, you can load and run it just like any other SavedModel:

1import numpy as np
2
3# Load the TensorRT-optimized model
4trt_model = tf.saved_model.load('my_trt_model')
5infer = trt_model.signatures['serving_default']
6
7# Create sample input
8sample_input = tf.constant(np.random.rand(1, 784).astype(np.float32))
9
10# Run inference
11output = infer(sample_input)
12print(output)

Notice that the TensorRT model exposes the same signatures as the original SavedModel, so you do not need to change your serving infrastructure.

INT8 Precision with Calibration

If you want to use INT8 precision for even faster inference, you need to provide a calibration dataset. TensorRT uses this data to determine the optimal quantization ranges for each layer:

1def calibration_input_fn():
2    for _ in range(100):
3        yield (np.random.rand(1, 784).astype(np.float32),)
4
5converter = trt.TrtGraphConverterV2(
6    input_saved_model_dir='my_saved_model',
7    precision_mode=trt.TrtPrecisionMode.INT8,
8    use_calibration=True
9)
10
11converter.convert(calibration_input_fn=calibration_input_fn)
12converter.save('my_trt_int8_model')

Use representative real data for calibration rather than random inputs. The quality of your calibration dataset directly affects the accuracy of the INT8 model.

Common Pitfalls

• Unsupported operations: Not all TensorFlow operations have TensorRT equivalents. When TF-TRT encounters an unsupported op, it falls back to TensorFlow for that portion of the graph, which can reduce the performance gain. Check the TensorRT compatibility matrix before converting.
• Shape mismatches at runtime: TensorRT engines are optimized for specific input shapes. If your inference inputs have dynamic shapes that differ from what was seen during conversion, you may encounter errors or degraded performance. Use the dynamic_shape_profile option to handle variable batch sizes.
• Ignoring precision validation: Switching from FP32 to FP16 or INT8 can introduce small numerical differences. Always compare the output of the converted model against the original on a representative test set to verify that accuracy remains acceptable for your use case.
• Insufficient workspace memory: Setting max_workspace_size_bytes too low prevents TensorRT from applying its most aggressive optimizations. Start with 1 GB and increase if you have available GPU memory. Conversely, setting it too high can cause out-of-memory errors on smaller GPUs.
• Version incompatibility: The TensorRT version must be compatible with both your TensorFlow version and the NVIDIA driver installed on your deployment machine. Mismatched versions are one of the most common causes of conversion failures. Always verify the compatibility matrix in the NVIDIA documentation.

Summary

• TensorRT optimizes TensorFlow SavedModels for faster GPU inference through layer fusion, kernel tuning, and precision reduction.
• Use TrtGraphConverterV2 from the tensorflow.python.compiler.tensorrt module to perform the conversion.
• Choose FP16 precision for a good speed-accuracy tradeoff, or INT8 with a calibration dataset for maximum throughput.
• The converted model retains the same SavedModel signatures, making deployment straightforward.
• Always validate the converted model's accuracy against the original before deploying to production.
• Ensure version compatibility between TensorFlow, TensorRT, and NVIDIA drivers on your target hardware.