Python tensorrt_rewriter_config Examples

Example #1

File: tf_trt_integration_test_base.py Project: zzhangncsu/tensorflow

    def _GetConfigProto(self, run_params, graph_state):
        """Get config proto based on specific settings."""
        if graph_state != GraphState.ORIGINAL and run_params.use_optimizer:
            conversion_params = self.GetConversionParams(run_params)
            rewriter_cfg = trt_convert.tensorrt_rewriter_config(
                conversion_params.rewriter_config,
                conversion_params.max_batch_size,
                conversion_params.max_workspace_size_bytes,
                conversion_params.precision_mode,
                conversion_params.minimum_segment_size,
                conversion_params.is_dynamic_op,
                conversion_params.maximum_cached_engines,
                conversion_params.cached_engine_batch_sizes)

            graph_options = config_pb2.GraphOptions(
                rewrite_options=rewriter_cfg)
        else:
            graph_options = config_pb2.GraphOptions()

        gpu_options = config_pb2.GPUOptions()
        gpu_options.allow_growth = True
        if trt_convert.get_linked_tensorrt_version()[0] == 3:
            gpu_options.per_process_gpu_memory_fraction = 0.50

        config = config_pb2.ConfigProto(gpu_options=gpu_options,
                                        graph_options=graph_options)
        return config

Example #2

File: trt_convert_test.py Project: yuchenlin/tensorflow

 def testTensorrtRewriterConfig(self):
     """Test case for trt_convert.tensorrt_rewriter_config()."""
     rewriter_cfg = trt_convert.tensorrt_rewriter_config(
         max_batch_size=128,
         max_workspace_size_bytes=1234,
         precision_mode="INT8",
         minimum_segment_size=10,
         is_dynamic_op=True,
         maximum_cached_engines=2,
         cached_engine_batch_sizes=[1, 128])
     trt_optimizer = None
     for optimizer in rewriter_cfg.custom_optimizers:
         if optimizer.name == "TensorRTOptimizer":
             self.assertTrue(trt_optimizer is None)
             trt_optimizer = optimizer
     self.assertTrue(trt_optimizer is not None)
     for key in [
             "minimum_segment_size", "max_batch_size", "is_dynamic_op",
             "max_workspace_size_bytes", "precision_mode",
             "maximum_cached_engines", "cached_engine_batches"
     ]:
         self.assertTrue(key in trt_optimizer.parameter_map)
     self.assertEqual(10,
                      trt_optimizer.parameter_map["minimum_segment_size"].i)
     self.assertEqual(128, trt_optimizer.parameter_map["max_batch_size"].i)
     self.assertEqual(True, trt_optimizer.parameter_map["is_dynamic_op"].b)
     self.assertEqual(
         1234, trt_optimizer.parameter_map["max_workspace_size_bytes"].i)
     self.assertEqual(trt_convert._to_bytes("INT8"),
                      trt_optimizer.parameter_map["precision_mode"].s)
     self.assertEqual(
         2, trt_optimizer.parameter_map["maximum_cached_engines"].i)
     self.assertEqual(
         [1, 128],
         trt_optimizer.parameter_map["cached_engine_batches"].list.i)

Example #3

File: trt_convert_test.py Project: ThunderQi/tensorflow

 def testTensorrtRewriterConfig(self):
   """Test case for trt_convert.tensorrt_rewriter_config()."""
   rewriter_cfg = trt_convert.tensorrt_rewriter_config(
       max_batch_size=128,
       max_workspace_size_bytes=1234,
       precision_mode="INT8",
       minimum_segment_size=10,
       is_dynamic_op=True,
       maximum_cached_engines=2,
       cached_engine_batch_sizes=[1, 128])
   trt_optimizer = None
   for optimizer in rewriter_cfg.custom_optimizers:
     if optimizer.name == "TensorRTOptimizer":
       self.assertTrue(trt_optimizer is None)
       trt_optimizer = optimizer
   self.assertTrue(trt_optimizer is not None)
   for key in [
       "minimum_segment_size", "max_batch_size", "is_dynamic_op",
       "max_workspace_size_bytes", "precision_mode", "maximum_cached_engines",
       "cached_engine_batches"
   ]:
     self.assertTrue(key in trt_optimizer.parameter_map)
   self.assertEqual(10, trt_optimizer.parameter_map["minimum_segment_size"].i)
   self.assertEqual(128, trt_optimizer.parameter_map["max_batch_size"].i)
   self.assertEqual(True, trt_optimizer.parameter_map["is_dynamic_op"].b)
   self.assertEqual(1234,
                    trt_optimizer.parameter_map["max_workspace_size_bytes"].i)
   self.assertEqual(
       trt_convert._to_bytes("INT8"),
       trt_optimizer.parameter_map["precision_mode"].s)
   self.assertEqual(2, trt_optimizer.parameter_map["maximum_cached_engines"].i)
   self.assertEqual(
       [1, 128],
       trt_optimizer.parameter_map["cached_engine_batches"].list.i)

Example #4

File: tf_trt_integration_test_base.py Project: abhinav-upadhyay/tensorflow

  def _GetConfigProto(self, run_params, graph_state):
    """Get config proto based on specific settings."""
    if graph_state != GraphState.ORIGINAL and run_params.use_optimizer:
      conversion_params = self.GetConversionParams(run_params)
      rewriter_cfg = trt_convert.tensorrt_rewriter_config(
          conversion_params.rewriter_config, conversion_params.max_batch_size,
          conversion_params.max_workspace_size_bytes,
          conversion_params.precision_mode,
          conversion_params.minimum_segment_size,
          conversion_params.is_dynamic_op,
          conversion_params.maximum_cached_engines,
          conversion_params.cached_engine_batch_sizes)

      graph_options = config_pb2.GraphOptions(rewrite_options=rewriter_cfg)
    else:
      graph_options = config_pb2.GraphOptions()

    gpu_options = config_pb2.GPUOptions()
    gpu_options.allow_growth = True
    if trt_convert.get_linked_tensorrt_version()[0] == 3:
      gpu_options.per_process_gpu_memory_fraction = 0.50

    config = config_pb2.ConfigProto(
        gpu_options=gpu_options, graph_options=graph_options)
    return config