def get_tf_graph_from_graph_model_v2()

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

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

    def get_tf_graph_from_graph_model_v2(self):
        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))
            break
        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)
        tf_graph = frozen_func.graph.as_graph_def(add_shapes=True)
        return tf_graph