There’s an existing thread on this: https://github.com/tensorflow/tensorflow/issues/53529.
Following the suggestions, I spun up a Workbench instance on GCP with TensorFlow 2.6 Enterprise which comes with TensorRT:
>>> tf.sysconfig.get_build_info()
OrderedDict([('cpu_compiler', '/usr/bin/gcc-5'),
('cuda_compute_capabilities',
['compute_37',
'compute_60',
'compute_61',
'compute_70',
'compute_75',
'compute_80']),
('cuda_version', '11.0'),
('cudnn_version', '8'),
('is_cuda_build', True),
('is_rocm_build', False),
('is_tensorrt_build', True)])
I tried the following code:
import tensorflow as tf
resnet = tf.keras.applications.ResNet50(weights="imagenet", include_top=True)
resnet.save("resnet")
TENSORRT_MODEL_DIR = f"tensorrt-resnet"
params = tf.experimental.tensorrt.ConversionParams(
precision_mode='FP16'
)
converter = tf.experimental.tensorrt.Converter(
input_saved_model_dir="resnet",
conversion_params=params
)
converter.convert()
converter.save(TENSORRT_MODEL_DIR)
The kernel restarts automatically when it tries to create the Converter
object. I have tried with tensorflow.python.compiler.tensorrt
too (as shown here) but it didn’t help.
Any help?