leimao/Frozen-Graph-TensorFlow

Issue with passing of the arguments?

Zain-Gill123 opened this issue · 0 comments

Hey! This might come out as a very newbie question, I'm very new to this all.
I am having issues with what to pass as the arguments of the save() and save_as_pb() fucntions.

Code Snippet:

!mkdir -p saved_model # Making he directory to save the saved model
model = tf.train.load_checkpoint('/content/test1/shanghaitech') # Loading the pre-trained checkpoint as model

# Your Code for the save() and the save_as_pb() functions:
from tensorflow.python.tools import freeze_graph

def save(self, directory, filename):

    if not os.path.exists(directory):
        os.makedirs(directory)
    filepath = os.path.join(directory, filename + '.ckpt')
    self.saver.save(self.sess, filepath)
    return filepath

def save_as_pb(self, directory, filename):

    if not os.path.exists(directory):
        os.makedirs(directory)

    # Save check point for graph frozen later
    ckpt_filepath = self.save(directory=directory, filename=filename)
    pbtxt_filename = filename + '.pbtxt'
    pbtxt_filepath = os.path.join(directory, pbtxt_filename)
    pb_filepath = os.path.join(directory, filename + '.pb')
    # This will only save the graph but the variables will not be saved.
    # You have to freeze your model first.
    tf.train.write_graph(graph_or_graph_def=self.sess.graph_def, logdir=directory, name=pbtxt_filename, as_text=True)

    # Freeze graph
    # Method 1
    freeze_graph.freeze_graph(input_graph=pbtxt_filepath, input_saver='', input_binary=False, input_checkpoint=ckpt_filepath, output_node_names='cnn/output', restore_op_name='save/restore_all', filename_tensor_name='save/Const:0', output_graph=pb_filepath, clear_devices=True, initializer_nodes='')
    
    return pb_filepath

# Calling the Function save_as_pb():
file_path = save_as_pb(self, directory='/content/test1/', filename='shanghaitech')
print(file_path)

The Issue:

When I call the save_as_pb() fucntion, I don't know what to pass to the "self" argument. I have tried passing "model" in place of "self" but of no use.

Any Kind of guidance would be highly appreciated!
Cheers
-Zain Gill