def __extract_input_and_output_tensors_from_saved_model_v2()

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

• 19 lines of code
• 6 McCabe index (conditional complexity)

    def __extract_input_and_output_tensors_from_saved_model_v2(self) -> None:
        from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2

        tags = self.__get_tag_set()
        loaded = tf.compat.v2.saved_model.load(self.model_path.as_posix(), tags=tags)
        for shape in self.__data_shape.values():
            tensor_spec = tf.TensorSpec(tuple(shape))
        if len(loaded.signatures) == 0:
            f = loaded.__call__.get_concrete_function(tensor_spec)
        elif 'serving_default' in loaded.signatures:
            f = loaded.signatures['serving_default']
        else:
            f = loaded.signatures[list(loaded.signatures.keys())[0]]
        frozen_func = convert_variables_to_constants_v2(f, lower_control_flow=True)

        for tensor in frozen_func.inputs:
            self.__input_tensor_names.append(tensor.name)
        for tensor in frozen_func.outputs:
            self.__output_tensor_names.append(tensor.name)

        self.__input_tensors = frozen_func.inputs
        self.__output_tensors = frozen_func.outputs