def __call__(self):
"""
Writes out artifacts from a TensorFlow Graph.
Returns:
Tuple[tf.Graph, Sequence[str]]: The TensorFlow graph, and the names of its outputs.
"""
(graph, outputs), _ = misc.try_call(self._graph)
misc.lazy_write(
contents=lambda: graph.as_graph_def().SerializeToString(),
path=self.path)
if self.tensorboard_dir:
G_LOGGER.info("Writing tensorboard events to {:}".format(
self.tensorboard_dir))
train_writer = tf.compat.v1.summary.FileWriter(
self.tensorboard_dir)
train_writer.add_graph(graph)
if self.engine_dir is not None:
graphdef = graph.as_graph_def()
segment_number = 0
for node in graphdef.node:
if node.op == "TRTEngineOp":
engine = node.attr["serialized_segment"].s
if self.engine_dir is not None:
misc.lazy_write(
contents=engine,
path=os.path.join(
self.engine_dir,
"segment-{:}".format(segment_number)))
segment_number += 1
return graph, outputs
def infer_impl(self, feed_dict):
G_LOGGER.extra_verbose("Received feed_dict: {:}".format(feed_dict))
start = time.time()
inference_outputs = self.sess.run(self.output_names,
feed_dict=feed_dict,
options=self.run_options,
run_metadata=self.run_metadata)
end = time.time()
out_dict = OrderedDict()
for name, out in zip(self.output_names, inference_outputs):
out_dict[name] = out
self.inference_time = end - start
def generate_timeline():
from tensorflow.python.client import timeline
t1 = timeline.Timeline(self.run_metadata.step_stats)
return t1.generate_chrome_trace_format()
if self.timeline_dir is not None:
misc.lazy_write(contents=generate_timeline,
path=os.path.join(
self.timeline_dir,
"run-{:}".format(self.num_inferences)),
mode="w")
self.num_inferences += 1
return out_dict
def __call__(self):
"""
Saves an engine to the provided path.
Returns:
trt.ICudaEngine: The engine that was saved.
"""
engine, _ = misc.try_call(self._engine)
misc.lazy_write(contents=lambda: engine.serialize(), path=self.path)
return engine
def __call__(self):
"""
Saves an ONNX model to the specified path.
Returns:
onnx.ModelProto: The model, after saving it.
"""
model, _ = misc.try_call(self._model)
misc.lazy_write(contents=lambda: model.SerializeToString(), path=self.path)
return model
def __call__(self):
"""
Saves an engine to the provided path.
Returns:
trt.ICudaEngine: The engine that was saved.
"""
engine, owns_engine = misc.try_call(self._engine)
with contextlib.ExitStack() as stack:
if owns_engine:
stack.enter_context(misc.FreeOnException([engine]))
misc.lazy_write(contents=lambda: engine.serialize(),
path=self.path)
return engine