TensorRT profile generation mode - #36729

This PR adds a profile generation to dynamic shapes mode.

This features completes the ground work for dynamic shapes mode. Profile generation mode is enabled during build(), and collects information of the input shapes for TRT engines. This information is used to generate a singe engine with multiple profiles.

Dynamic shape mode is an experimental feature. The advantage of dynamic shapes is that we can use a single TensorRT engine with multiple profiles to handle input data with different shapes. This saves memory: all input profiles for an engine can share the same weights. Input dimensions that do not influence the number of model parameters can be dynamic (batch size, sequence length, even image dimensions for fully convolutional networks). Currently we generate a separate profile for each input shape that was provided during build mode.

This feature does not support defining profiles during runtime, only during build mode. This makes it important for users to feed data to build() with all interesting shapes so that the single engine can execute those shapes. Currently this feature does not support INT8 calibration.

This feature is not supposed to change any of the existing behavior of TF-TRT in implicit batch mode. Explicit batch mode is changed: it is now required to use build() to generate the engines.

It is expected to use this feature for elementwise ops. We are soon going to enable all the existing converters to work with this feature. An example on how to use this feature:

    root = _model()
    save.save(root, input_saved_model_dir,
              {signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: root.run})

    rewriter_config_with_trt = rewriter_config_pb2.RewriterConfig()
    rewriter_config_with_trt.optimizers.extend(
        ["constfold", "layout", "constfold"])
    rewriter_config_with_trt.meta_optimizer_iterations = (
        rewriter_config_pb2.RewriterConfig.ONE)
    optimizer = rewriter_config_with_trt.custom_optimizers.add()
    rewriter_config_with_trt.custom_optimizers.add().name = "constfold"
    optimizer.name = "TensorRTOptimizer"
    optimizer.parameter_map["minimum_segment_size"].i = 1
    optimizer.parameter_map["is_dynamic_op"].b = True
    optimizer.parameter_map["maximum_cached_engines"].i = 1
    optimizer.parameter_map["use_implicit_batch"].b = False

    conversion_params = trt_convert.DEFAULT_TRT_CONVERSION_PARAMS._replace(
        rewriter_config_template=rewriter_config_with_trt)
    converter = trt_convert.TrtGraphConverterV2(input_saved_model_dir=input_saved_model_dir,
                                    conversion_params=conversion_params)
    converter.convert()

    input_shapes = [(2, 4), (3, 9), (64, 128)]
    def my_input_fn():
      for x in input_shapes:
        yield (np.random.normal(size=x).astype(np.float32),)

    converter.build(input_fn=my_input_fn)

    output_saved_model_dir = self.mkdtemp()
    converter.save(output_saved_model_dir=output_saved_model_dir)