-
Notifications
You must be signed in to change notification settings - Fork 16
Expand file tree
/
Copy pathtfconverter.py
More file actions
22 lines (18 loc) · 785 Bytes
/
tfconverter.py
File metadata and controls
22 lines (18 loc) · 785 Bytes
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
import tensorflow as tf
from google.protobuf import text_format
from tensorflow.python.platform import gfile
def pbtxt_to_graphdef(filename):
with open(filename, 'r') as f:
graph_def = tf.GraphDef()
file_content = f.read()
text_format.Merge(file_content, graph_def)
tf.import_graph_def(graph_def, name='')
tf.train.write_graph(graph_def, './', 'protobuf.pb', as_text=False)
def graphdef_to_pbtxt(filename):
with gfile.FastGFile(filename,'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
tf.import_graph_def(graph_def, name='')
tf.train.write_graph(graph_def, './', filename+'txt', as_text=True)
return
graphdef_to_pbtxt('model/frozen_inference_graph.pb') # here you can write the name of the file to be converted