A Python package for seamlessly converting Efficientnet TensorFlow models to ONNX and TensorRT formats, enabling optimized inference on NVIDIA GPUs.
- TensorFlow to ONNX Conversion: Convert SavedModel or model weights to ONNX format with FP32 or FP16 precision
- ONNX to TensorRT Conversion: Build optimized TensorRT engines with dynamic batch sizes
- Direct TF to TRT Pipeline: One-step conversion from TensorFlow to TensorRT
- INT8 Quantization Support: Improve inference speed with INT8 calibration
- Comprehensive Logging: Detailed conversion process tracking and validation
Before installing the package, ensure you have the following:
- Python 3.8 or higher - Check with
python --version - pip (Python package manager) - Check with
pip --version - NVIDIA GPU with CUDA support
- CUDA Toolkit 11.7+ - Download from NVIDIA
- TensorRT 8.5+ - Download from NVIDIA
- Download the wheel file provided to you
- Open a terminal and navigate to the directory containing the .whl file.
- Install the package using pip.
- If dependencies are not automatically installed with the wheel file, you can install them using the provided
requirements.txtfile:
- Install the wheel and dependencies, then import the converter:
from quantizeeffinet import ModelConverter - Create a converter instance:
converter = ModelConverter() - Convert TensorFlow to ONNX or straight to TensorRT, with optional FP16/INT8 settings and dynamic batching.
ModelConverter provides three utilities: convert TensorFlow models to ONNX, build TensorRT engines from ONNX, and run a one-call TensorFlow to TensorRT pipeline for deployment-ready inference engines. These three utilities are accessed by three functions.
- The original input shape (NHWC) will be changed to (NCHW)
- All the functions'
output_model_pathcould beNonethis way the functions return the created models. However, it is not recommended to use these models for inference right away. The best practice is to restart the Python kernel, thus freeing up all allocated memory. After that, with a new run, you could run inference without any errors. - Fp16 TensorRT engines are the bes choice for now, as they provide the fastes inference wihout any accuracy loss.
- Converts a
tf.kerasSavedModel or a.h5Weights file to an ONNX model. - I you pass a
.h5file that only contains weights, you should also specify the base model, like EfficientNetB3, B5, B6 - The keras model input shape must be
(None, 224, 224, 3). - Supports exporting to FP32 or FP16 ONNX graph.
- You do not have to change opset argument in the function call
opset=13is sufficient.
from quantizeeffinet import ModelConverter
converter = ModelConverter()
onnx_model = converter.tf_to_onnx(
input_model="path/to/efficientnet_b3.h5",
output_path="path/to/output/model.onnx",
precision="fp16",
only_weigths_of_model='EfficientNetB3'
)
- Parses an ONNX model and builds a TensorRT engine.
- The input model could be passed as a path to a
model.onnxfile or as aonnx.ModelProtoobject. - The onnx model input shape must be
(-1, 3, 224, 224). - If no path to save the engine and
auto_generate_engine_path=Trueit auto-generates the path. - Supports exporting to FP32, FP16 or Int8 trt engines.
- TensorRT uses dynamic batching. This means that you have to specify 3 arguments:
min_batch: int = 1, this is the minimum nuber of images that a batch could ever contain (Should stay 1).opt_batch: int = 32, this is the number of images that most of the inference batches will contain (Should be the same as max_batch).max_batch: int = 32, this is the maximum nuber of images that a batch could ever contain (Should be the same as max_batch).
from quantizeeffinet import ModelConverter
converter = ModelConverter()
engine = converter.onnx_to_trt(
input_model="path/to/model.onnx",
engine_file_path="path/to/output/model.engine",
precision="fp16",
min_batch=1,
opt_batch=16,
max_batch=32
)
- INT8 requires a calibration step or an existing cache.
- INT8 mode supports using a directory of images, a single image, a list of images, or an existing calibration cache to avoid recalibration.
- If no cache value is given then images are going to be used and no cache will be saved, if no images are given then cache value will be used. If both are pased to the function then the cache will be used for faster conversion.
from quantizeeffinet import ModelConverter
converter = ModelConverter()
model=converter.onnx_to_trt(
input_model=path/to/model.onnx,
engine_file_path=path/to/int8_efficientnet.engine,
min_batch=1,
max_batch=32,
opt_batch=32,
precision="int8",
calibration_images=path/to/calibration/images/dir,
calibration_cache=path/to/calibration.cache,
)
Runs the end-to-end pipeline in one call. Exports a TensorFlow model to ONNX, then builds a TensorRT engine with the selected precision and batch profiles, streamlining deployment.
- If you pass a
.h5file that only contains weights, you should also specify the base model, like EfficientNetB3, B5, B6 - When using INT8, provide calibration images or a calibration cache.
from quantizeeffinet import ModelConverter
converter = ModelConverter()
model=converter.tf_to_trt(
input_model="path/to/efficientnet_b3.h5",
engine_file_path=path/to/int8_efficientnet.engine,
only_weigths_of_model='EfficientNetB3',
min_batch=1,
opt_batch=32,
max_batch=32,
precision="int8",
calibration_images=path/to/calibration/images/dir,
calibration_cache=path/to/calibration.cache,
)