def testTrtGraphConverter_OnlineConversion(self, device):
"""Test case for TF-TRT conversion using Grappler directly."""
if not is_tensorrt_enabled():
return
conversion_params = trt_convert.DEFAULT_TRT_CONVERSION_PARAMS._replace(
precision_mode=trt_convert.TrtPrecisionMode.FP32, is_dynamic_op=True)
config = self._GetConfigProto(
rewriter_config=trt_convert.get_tensorrt_rewriter_config(
conversion_params, is_v2=False))
with ops.Graph().as_default():
# Online conversion requires a frozen graph, so we reuse inp1 as the var
# argument.
inp1 = array_ops.placeholder(
dtype=dtypes.float32, shape=[None, 1, 1], name="input1")
inp2 = array_ops.placeholder(
dtype=dtypes.float32, shape=[None, 1, 1], name="input2")
if device:
with ops.device(device):
TrtConvertTest._GetGraph(inp1, inp2, inp1)
else:
TrtConvertTest._GetGraph(inp1, inp2, inp1)
with self.session(config=config) as sess:
self._TestRun(sess, batch_size=1)
def _GetConfigProto(self, run_params, graph_state):
"""Get config proto based on specific settings."""
conversion_params = self.GetConversionParams(run_params)
if graph_state != GraphState.ORIGINAL and run_params.use_optimizer:
rewriter_cfg = trt_convert.get_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_batches,
conversion_params.use_calibration)
graph_options = config_pb2.GraphOptions(
rewrite_options=rewriter_cfg)
else:
graph_options = config_pb2.GraphOptions()
if conversion_params.rewriter_config is not None:
graph_options.rewrite_options.CopyFrom(
conversion_params.rewriter_config)
config = config_pb2.ConfigProto(gpu_options=self._GetGPUOptions(),
graph_options=graph_options)
return config
def GetTrtRewriterConfig(self,
run_params,
conversion_params,
disable_non_trt_optimizers=False):
return trt_convert.get_tensorrt_rewriter_config(
conversion_params=conversion_params,
is_v2=run_params.is_v2,
disable_non_trt_optimizers=disable_non_trt_optimizers)
def testGetTensorrtRewriterConfigTemplate(self):
"""Test case for TrtGraphConverter.get_tensorrt_rewriter_config()."""
if not is_tensorrt_enabled():
return
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 = 10
optimizer.parameter_map["max_batch_size"].i = 128
optimizer.parameter_map["is_dynamic_op"].b = True
optimizer.parameter_map["max_workspace_size_bytes"].i = 1234
optimizer.parameter_map["precision_mode"].s = trt_convert._to_bytes(
trt_convert.TrtPrecisionMode.INT8)
optimizer.parameter_map["maximum_cached_engines"].i = 2
optimizer.parameter_map["use_calibration"].b = False
optimizer.parameter_map["use_implicit_batch"].b = True
conversion_params = trt_convert.DEFAULT_TRT_CONVERSION_PARAMS._replace(
rewriter_config_template=rewriter_config_with_trt)
rewriter_cfg = trt_convert.get_tensorrt_rewriter_config(
conversion_params=conversion_params)
self.assertEqual(["constfold", "layout", "constfold"],
rewriter_cfg.optimizers)
self.assertEqual(rewriter_config_pb2.RewriterConfig.ONE,
rewriter_cfg.meta_optimizer_iterations)
trt_optimizer = None
for optimizer in rewriter_cfg.custom_optimizers:
if optimizer.name == "TensorRTOptimizer":
self.assertIsNone(trt_optimizer)
trt_optimizer = optimizer
self.assertIsNotNone(trt_optimizer)
for key in [
"minimum_segment_size", "max_batch_size", "is_dynamic_op",
"max_workspace_size_bytes", "precision_mode",
"maximum_cached_engines"
]:
self.assertIn(key, 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(False,
trt_optimizer.parameter_map["use_calibration"].b)
self.assertEqual(True,
trt_optimizer.parameter_map["use_implicit_batch"].b)
def GetTrtRewriterConfig(self,
run_params,
conversion_params,
disable_non_trt_optimizers=False,
use_implicit_batch=True):
rewriter_config = trt_convert.get_tensorrt_rewriter_config(
conversion_params=conversion_params,
is_v2=run_params.is_v2,
disable_non_trt_optimizers=disable_non_trt_optimizers)
for optimizer in rewriter_config.custom_optimizers:
if optimizer.name == "TensorRTOptimizer":
optimizer.parameter_map["use_implicit_batch"].b = use_implicit_batch
return rewriter_config
def _GetConfigProto(self, run_params, graph_state):
"""Get config proto based on specific settings."""
conversion_params = self.GetConversionParams(run_params)
if graph_state == GraphState.INFERENCE and run_params.convert_online:
rewriter_cfg = trt_convert.get_tensorrt_rewriter_config(conversion_params)
graph_options = config_pb2.GraphOptions(rewrite_options=rewriter_cfg)
else:
graph_options = config_pb2.GraphOptions()
if conversion_params.rewriter_config_template is not None:
graph_options.rewrite_options.CopyFrom(
conversion_params.rewriter_config_template)
config = config_pb2.ConfigProto(
gpu_options=self._GetGPUOptions(), graph_options=graph_options)
return config
def _GetConfigProto(self, run_params, graph_state):
"""Get config proto based on specific settings."""
conversion_params = self.GetConversionParams(run_params)
if graph_state == GraphState.INFERENCE and run_params.use_optimizer:
rewriter_cfg = trt_convert.get_tensorrt_rewriter_config(conversion_params)
graph_options = config_pb2.GraphOptions(rewrite_options=rewriter_cfg)
else:
graph_options = config_pb2.GraphOptions()
if conversion_params.rewriter_config_template is not None:
graph_options.rewrite_options.CopyFrom(
conversion_params.rewriter_config_template)
config = config_pb2.ConfigProto(
gpu_options=self._GetGPUOptions(), graph_options=graph_options)
return config
def _GetConfigProto(self, run_params, graph_state):
"""Get config proto based on specific settings."""
conversion_params = self.GetConversionParams(run_params)
max_batch_size = self.GetMaxBatchSize(run_params)
if graph_state == GraphState.INFERENCE and run_params.convert_online:
rewriter_cfg = trt_convert.get_tensorrt_rewriter_config(
conversion_params,
is_dynamic_op=run_params.dynamic_engine,
max_batch_size=max_batch_size)
graph_options = config_pb2.GraphOptions(rewrite_options=rewriter_cfg)
else:
graph_options = config_pb2.GraphOptions()
config = config_pb2.ConfigProto(
gpu_options=self._GetGPUOptions(), graph_options=graph_options)
return config
def testGetTensorrtRewriterConfig(self):
"""Test case for TrtGraphConverter.get_tensorrt_rewriter_config()."""
if not is_tensorrt_enabled():
return
conversion_params = trt_convert.DEFAULT_TRT_CONVERSION_PARAMS._replace(
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_batches=[1, 128])
rewriter_cfg = trt_convert.get_tensorrt_rewriter_config(
conversion_params=conversion_params)
self.assertEqual(["constfold", "layout", "constfold"],
rewriter_cfg.optimizers)
self.assertEqual(rewriter_config_pb2.RewriterConfig.ONE,
rewriter_cfg.meta_optimizer_iterations)
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)
def testGetTensorrtRewriterConfig(self):
"""Test case for TrtGraphConverter.get_tensorrt_rewriter_config()."""
if not is_tensorrt_enabled():
return
conversion_params = trt_convert.DEFAULT_TRT_CONVERSION_PARAMS._replace(
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_batches=[1, 128])
rewriter_cfg = trt_convert.get_tensorrt_rewriter_config(
conversion_params=conversion_params)
self.assertEqual(["constfold", "layout", "constfold"],
rewriter_cfg.optimizers)
self.assertEqual(rewriter_config_pb2.RewriterConfig.ONE,
rewriter_cfg.meta_optimizer_iterations)
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)
def _GetConfigProto(self, run_params, graph_state):
"""Get config proto based on specific settings."""
conversion_params = self.GetConversionParams(run_params)
if graph_state != GraphState.ORIGINAL and run_params.use_optimizer:
rewriter_cfg = trt_convert.get_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_batches,
conversion_params.use_calibration)
graph_options = config_pb2.GraphOptions(rewrite_options=rewriter_cfg)
else:
graph_options = config_pb2.GraphOptions()
if conversion_params.rewriter_config is not None:
graph_options.rewrite_options.CopyFrom(
conversion_params.rewriter_config)
config = config_pb2.ConfigProto(
gpu_options=self._GetGPUOptions(), graph_options=graph_options)
return config
