What's the state of TensorRT for newer versions of TensorFlow?

General Discussion
, , ,
#1

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?