def testCreateInferenceGraph_DynamicOp(self):
if not trt_convert.is_tensorrt_enabled():
return
trt_convert.enable_test_value()
tmp_dir = self.get_temp_dir()
input_saved_model_dir = os.path.join(tmp_dir, "in_dir2")
output_saved_model_dir = os.path.join(tmp_dir, "out_dir2")
self._WriteInputSavedModel(input_saved_model_dir)
output_graph_def = trt_convert.create_inference_graph(
None,
None,
is_dynamic_op=True,
maximum_cached_engines=2,
input_saved_model_dir=input_saved_model_dir,
output_saved_model_dir=output_saved_model_dir,
session_config=self._GetConfigProto())
# Test the output GraphDef.
with ops.Graph().as_default():
importer.import_graph_def(output_graph_def, name="")
with self.test_session(config=self._GetConfigProto()) as sess:
# Run with batch size 1, a new engine is created and cached.
self._TestRun(sess, 1, True)
# Run with batch size 2, a new engine is created and cached.
self._TestRun(sess, 2, True)
# Run with batch size 3, since the number of cached engines has reached
# the max, it should fall back to TF function.
self._TestRun(sess, 3, False)
# Test the output SavedModel
with ops.Graph().as_default():
with self.test_session(config=self._GetConfigProto()) as sess:
loader.load(sess, [tag_constants.SERVING],
output_saved_model_dir)
# Run with batch size 1, a new engine is created and cached.
self._TestRun(sess, 1, True)
# Run with batch size 2, a new engine is created and cached.
self._TestRun(sess, 2, True)
# Run with batch size 3, since the number of cached engines has reached
# the max, it should fall back to TF function.
self._TestRun(sess, 3, False)
def testCreateInferenceGraph_DynamicOp(self):
if not trt_convert.is_tensorrt_enabled():
return
trt_convert.enable_test_value()
tmp_dir = self.get_temp_dir()
input_saved_model_dir = os.path.join(tmp_dir, "in_dir2")
output_saved_model_dir = os.path.join(tmp_dir, "out_dir2")
self._WriteInputSavedModel(input_saved_model_dir)
output_graph_def = trt_convert.create_inference_graph(
None,
None,
is_dynamic_op=True,
maximum_cached_engines=2,
input_saved_model_dir=input_saved_model_dir,
output_saved_model_dir=output_saved_model_dir,
session_config=self._GetConfigProto())
# Test the output GraphDef.
with ops.Graph().as_default():
importer.import_graph_def(output_graph_def, name="")
with self.test_session(config=self._GetConfigProto()) as sess:
# Run with batch size 1, a new engine is created and cached.
self._TestRun(sess, 1, True)
# Run with batch size 2, a new engine is created and cached.
self._TestRun(sess, 2, True)
# Run with batch size 3, since the number of cached engines has reached
# the max, it should fall back to TF function.
self._TestRun(sess, 3, False)
# Test the output SavedModel
with ops.Graph().as_default():
with self.test_session(config=self._GetConfigProto()) as sess:
loader.load(sess, [tag_constants.SERVING], output_saved_model_dir)
# Run with batch size 1, a new engine is created and cached.
self._TestRun(sess, 1, True)
# Run with batch size 2, a new engine is created and cached.
self._TestRun(sess, 2, True)
# Run with batch size 3, since the number of cached engines has reached
# the max, it should fall back to TF function.
self._TestRun(sess, 3, False)
def testCreateInferenceGraph_StaticOp(self):
if not trt_convert.is_tensorrt_enabled():
return
trt_convert.enable_test_value()
tmp_dir = self.get_temp_dir()
input_saved_model_dir = os.path.join(tmp_dir, "in_dir3")
output_saved_model_dir = os.path.join(tmp_dir, "out_dir3")
self._WriteInputSavedModel(input_saved_model_dir)
output_graph_def = trt_convert.create_inference_graph(
None,
None,
max_batch_size=1,
is_dynamic_op=False,
maximum_cached_engines=2, # This is noop, added just for testing.
input_saved_model_dir=input_saved_model_dir,
output_saved_model_dir=output_saved_model_dir,
session_config=self._GetConfigProto())
# Test the output GraphDef.
with ops.Graph().as_default():
importer.import_graph_def(output_graph_def, name="")
with self.test_session(config=self._GetConfigProto()) as sess:
# Run with batch size 1, the default engine embedded in the graphdef
# will be used.
self._TestRun(sess, 1, True)
# Run with batch size 2, which exceed the max_batch_size, it should fall
# back to TF function.
self._TestRun(sess, 2, False)
# Test the output SavedModel
with ops.Graph().as_default():
with self.test_session(config=self._GetConfigProto()) as sess:
loader.load(sess, [tag_constants.SERVING],
output_saved_model_dir)
# Run with batch size 1, the default engine embedded in the graphdef
# will be used.
self._TestRun(sess, 1, True)
# Run with batch size 2, which exceed the max_batch_size, it should fall
# back to TF function.
self._TestRun(sess, 2, False)
def testCreateInferenceGraph_StaticOp(self):
if not trt_convert.is_tensorrt_enabled():
return
trt_convert.enable_test_value()
tmp_dir = self.get_temp_dir()
input_saved_model_dir = os.path.join(tmp_dir, "in_dir3")
output_saved_model_dir = os.path.join(tmp_dir, "out_dir3")
self._WriteInputSavedModel(input_saved_model_dir)
output_graph_def = trt_convert.create_inference_graph(
None,
None,
max_batch_size=1,
is_dynamic_op=False,
maximum_cached_engines=2, # This is noop, added just for testing.
input_saved_model_dir=input_saved_model_dir,
output_saved_model_dir=output_saved_model_dir,
session_config=self._GetConfigProto())
# Test the output GraphDef.
with ops.Graph().as_default():
importer.import_graph_def(output_graph_def, name="")
with self.test_session(config=self._GetConfigProto()) as sess:
# Run with batch size 1, the default engine embedded in the graphdef
# will be used.
self._TestRun(sess, 1, True)
# Run with batch size 2, which exceed the max_batch_size, it should fall
# back to TF function.
self._TestRun(sess, 2, False)
# Test the output SavedModel
with ops.Graph().as_default():
with self.test_session(config=self._GetConfigProto()) as sess:
loader.load(sess, [tag_constants.SERVING], output_saved_model_dir)
# Run with batch size 1, the default engine embedded in the graphdef
# will be used.
self._TestRun(sess, 1, True)
# Run with batch size 2, which exceed the max_batch_size, it should fall
# back to TF function.
self._TestRun(sess, 2, False)
def setUpClass(cls):
"""Setup method for the module."""
super(TfTrtIntegrationTestBase, cls).setUpClass()
trt_convert.enable_test_value()
def setUpClass(cls): """Setup method for the module.""" super(TfTrtIntegrationTestBase, cls).setUpClass() trt_convert.enable_test_value()
