Spaces:
Sleeping
Sleeping
| from tensorflow.python.compiler.tensorrt import trt_convert as trt | |
| def convert_to_trt(input_model_path: str, trt_model_path: str) -> None: | |
| """Utility to convert and save an input SavedModel to an optimized TensorRT graph. | |
| Args: | |
| input_model_path: Path to the SavedModel to optimize. | |
| trt_model_path: Path to save the converted TensorRT graph. | |
| """ | |
| converter = trt.TrtGraphConverterV2( | |
| input_saved_model_dir=input_model_path, | |
| precision_mode=trt.TrtPrecisionMode.FP32, | |
| max_workspace_size_bytes=8000000000, | |
| ) | |
| converter.convert() | |
| converter.save(output_saved_model_dir=trt_model_path) | |
| print("Done Converting to TF-TRT FP32") | |