def __extract_input_and_output_tensors_from_saved_model()

in src/neo_loader/helpers/tf_model_helper.py [0:0]

• 17 lines of code
• 4 McCabe index (conditional complexity)

    def __extract_input_and_output_tensors_from_saved_model(self) -> None:
        # https://github.com/apache/incubator-tvm/blob/master/python/tvm/relay/frontend/tensorflow_parser.py#L73
        tf.reset_default_graph()

        tags = self.__get_tag_set()
        input_tensors = OrderedDict()
        output_tensors = OrderedDict()

        with tf.Session() as sess:
            meta_graph_def = tf.saved_model.loader.load(sess, tags, self.model_path.as_posix())
            for sig_def in meta_graph_def.signature_def.values():
                for input_tensor in sig_def.inputs.values():
                    input_tensors[input_tensor.name] = tf.get_default_graph().get_tensor_by_name(input_tensor.name)
                for output_tensor in sig_def.outputs.values():
                    output_tensors[output_tensor.name] = tf.get_default_graph().get_tensor_by_name(output_tensor.name)

        tf.reset_default_graph()

        self.__input_tensor_names = list(input_tensors.keys())
        self.__output_tensor_names = list(output_tensors.keys())
        self.__input_tensors = list(input_tensors.values())
        self.__output_tensors = list(output_tensors.values())