Commit 2b24e23f by Alexander Pivovarov Committed by Tianqi Chen

Simplify TF get_output_names (#3025)

parent 7ee5cbac
......@@ -85,11 +85,9 @@ class TFParser(object):
tags,
self._model_dir)
output_names = set()
for k in meta_graph_def.signature_def.keys():
outputs_tensor_info = meta_graph_def.signature_def[k].outputs
for output_tensor in outputs_tensor_info.values():
output_names.add(output_tensor.name)
output_names = [i.replace(":0", "") for i in output_names]
for sig_def in meta_graph_def.signature_def.values():
for output_tensor in sig_def.outputs.values():
output_names.add(output_tensor.name.replace(":0", ""))
return ",".join(output_names)
def _load_saved_model(self):
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment