convert_from_tensorflow.py 25 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607
  1. # Copyright (c) 2019 Guo Yejun
  2. #
  3. # This file is part of FFmpeg.
  4. #
  5. # FFmpeg is free software; you can redistribute it and/or
  6. # modify it under the terms of the GNU Lesser General Public
  7. # License as published by the Free Software Foundation; either
  8. # version 2.1 of the License, or (at your option) any later version.
  9. #
  10. # FFmpeg is distributed in the hope that it will be useful,
  11. # but WITHOUT ANY WARRANTY; without even the implied warranty of
  12. # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
  13. # Lesser General Public License for more details.
  14. #
  15. # You should have received a copy of the GNU Lesser General Public
  16. # License along with FFmpeg; if not, write to the Free Software
  17. # Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  18. # ==============================================================================
  19. import tensorflow as tf
  20. import numpy as np
  21. import sys, struct
  22. import convert_header as header
  23. __all__ = ['convert_from_tensorflow']
  24. class Operand(object):
  25. IOTYPE_INPUT = 1
  26. IOTYPE_OUTPUT = 2
  27. IOTYPE_INTERMEDIATE = IOTYPE_INPUT | IOTYPE_OUTPUT
  28. DTYPE_FLOAT = 1
  29. DTYPE_UINT8 = 4
  30. index = 0
  31. def __init__(self, name, dtype, dims):
  32. self.name = name
  33. self.dtype = dtype
  34. self.dims = dims
  35. self.iotype = 0
  36. self.used_count = 0
  37. self.index = Operand.index
  38. Operand.index = Operand.index + 1
  39. self.iotype2str = {Operand.IOTYPE_INPUT: 'in', Operand.IOTYPE_OUTPUT: 'out', Operand.IOTYPE_INTERMEDIATE: 'inout'}
  40. self.dtype2str = {Operand.DTYPE_FLOAT: 'DT_FLOAT', Operand.DTYPE_UINT8: 'DT_UINT8'}
  41. def add_iotype(self, iotype):
  42. self.iotype = self.iotype | iotype
  43. if iotype == Operand.IOTYPE_INPUT:
  44. self.used_count = self.used_count + 1
  45. def __str__(self):
  46. return "{}: (name: {}, iotype: {}, dtype: {}, dims: {}, used_count: {})".format(self.index,
  47. self.name, self.iotype2str[self.iotype], self.dtype2str[self.dtype],
  48. self.dims, self.used_count)
  49. def __lt__(self, other):
  50. return self.index < other.index
  51. class TFConverter:
  52. def __init__(self, graph_def, nodes, outfile, dump4tb):
  53. self.graph_def = graph_def
  54. self.nodes = nodes
  55. self.outfile = outfile
  56. self.dump4tb = dump4tb
  57. self.layer_number = 0
  58. self.output_names = []
  59. self.name_node_dict = {}
  60. self.edges = {}
  61. self.conv_activations = {'Relu':0, 'Tanh':1, 'Sigmoid':2, 'None':3, 'LeakyRelu':4}
  62. self.conv_paddings = {'VALID':0, 'SAME':1}
  63. self.pool_paddings = {'VALID':0, 'SAME':1}
  64. self.converted_nodes = set()
  65. self.conv2d_scope_names = set()
  66. self.conv2d_scopename_inputname_dict = {}
  67. self.dense_scope_names = set()
  68. self.dense_scopename_inputname_dict = {}
  69. self.op2code = {'Conv2D':1, 'DepthToSpace':2, 'MirrorPad':3, 'Maximum':4,
  70. 'MathBinary':5, 'MathUnary':6, 'AvgPool':7, 'MatMul':8}
  71. self.mathbin2code = {'Sub':0, 'Add':1, 'Mul':2, 'RealDiv':3, 'Minimum':4, 'FloorMod':5}
  72. self.mathun2code = {'Abs':0, 'Sin':1, 'Cos':2, 'Tan':3, 'Asin':4,
  73. 'Acos':5, 'Atan':6, 'Sinh':7, 'Cosh':8, 'Tanh':9, 'Asinh':10,
  74. 'Acosh':11, 'Atanh':12, 'Ceil':13, 'Floor':14, 'Round':15,
  75. 'Exp':16}
  76. self.mirrorpad_mode = {'CONSTANT':0, 'REFLECT':1, 'SYMMETRIC':2}
  77. self.name_operand_dict = {}
  78. def add_operand(self, name, type):
  79. node = self.name_node_dict[name]
  80. if name not in self.name_operand_dict:
  81. dtype = node.attr['dtype'].type
  82. if dtype == 0:
  83. dtype = node.attr['T'].type
  84. dims = [-1,-1,-1,-1]
  85. if 'shape' in node.attr:
  86. dims[0] = node.attr['shape'].shape.dim[0].size
  87. dims[1] = node.attr['shape'].shape.dim[1].size
  88. dims[2] = node.attr['shape'].shape.dim[2].size
  89. dims[3] = node.attr['shape'].shape.dim[3].size
  90. operand = Operand(name, dtype, dims)
  91. self.name_operand_dict[name] = operand;
  92. self.name_operand_dict[name].add_iotype(type)
  93. return self.name_operand_dict[name].index
  94. def dump_for_tensorboard(self):
  95. graph = tf.get_default_graph()
  96. tf.import_graph_def(self.graph_def, name="")
  97. tf.summary.FileWriter('/tmp/graph', graph)
  98. print('graph saved, run "tensorboard --logdir=/tmp/graph" to see it')
  99. def get_conv2d_params(self, conv2d_scope_name):
  100. knode = self.name_node_dict[conv2d_scope_name + '/kernel']
  101. bnode = self.name_node_dict[conv2d_scope_name + '/bias']
  102. if conv2d_scope_name + '/dilation_rate' in self.name_node_dict:
  103. dnode = self.name_node_dict[conv2d_scope_name + '/dilation_rate']
  104. else:
  105. dnode = None
  106. # the BiasAdd name is possible be changed into the output name,
  107. # if activation is None, and BiasAdd.next is the last op which is Identity
  108. if conv2d_scope_name + '/BiasAdd' in self.edges:
  109. anode = self.edges[conv2d_scope_name + '/BiasAdd'][0]
  110. if anode.op not in self.conv_activations:
  111. anode = None
  112. else:
  113. anode = None
  114. return knode, bnode, dnode, anode
  115. def get_dense_params(self, dense_scope_name):
  116. knode = self.name_node_dict[dense_scope_name + '/kernel']
  117. bnode = self.name_node_dict.get(dense_scope_name + '/bias')
  118. # the BiasAdd name is possible be changed into the output name,
  119. # if activation is None, and BiasAdd.next is the last op which is Identity
  120. anode = None
  121. if bnode:
  122. if dense_scope_name + '/BiasAdd' in self.edges:
  123. anode = self.edges[dense_scope_name + '/BiasAdd'][0]
  124. if anode.op not in self.conv_activations:
  125. anode = None
  126. else:
  127. anode = None
  128. return knode, bnode, anode
  129. def dump_complex_conv2d_to_file(self, node, f):
  130. assert(node.op == 'Conv2D')
  131. self.layer_number = self.layer_number + 1
  132. self.converted_nodes.add(node.name)
  133. scope_name = TFConverter.get_scope_name(node.name)
  134. #knode for kernel, bnode for bias, dnode for dilation, anode for activation
  135. knode, bnode, dnode, anode = self.get_conv2d_params(scope_name)
  136. if dnode is not None:
  137. dilation = struct.unpack('i', dnode.attr['value'].tensor.tensor_content[0:4])[0]
  138. else:
  139. dilation = 1
  140. if anode is not None:
  141. activation = anode.op
  142. else:
  143. activation = 'None'
  144. padding = node.attr['padding'].s.decode("utf-8")
  145. # conv2d with dilation > 1 generates tens of nodes, not easy to parse them, so use this tricky method.
  146. if dilation > 1 and scope_name + '/stack' in self.name_node_dict:
  147. if self.name_node_dict[scope_name + '/stack'].op == "Const":
  148. padding = 'SAME'
  149. padding = self.conv_paddings[padding]
  150. ktensor = knode.attr['value'].tensor
  151. filter_height = ktensor.tensor_shape.dim[0].size
  152. filter_width = ktensor.tensor_shape.dim[1].size
  153. in_channels = ktensor.tensor_shape.dim[2].size
  154. out_channels = ktensor.tensor_shape.dim[3].size
  155. kernel = np.frombuffer(ktensor.tensor_content, dtype=np.float32)
  156. kernel = kernel.reshape(filter_height, filter_width, in_channels, out_channels)
  157. kernel = np.transpose(kernel, [3, 0, 1, 2])
  158. has_bias = 1
  159. np.array([self.op2code[node.op], dilation, padding, self.conv_activations[activation], in_channels, out_channels, filter_height, has_bias], dtype=np.uint32).tofile(f)
  160. kernel.tofile(f)
  161. btensor = bnode.attr['value'].tensor
  162. if btensor.tensor_shape.dim[0].size == 1:
  163. bias = struct.pack("f", btensor.float_val[0])
  164. else:
  165. bias = btensor.tensor_content
  166. f.write(bias)
  167. input_name = self.conv2d_scopename_inputname_dict[scope_name]
  168. input_operand_index = self.add_operand(input_name, Operand.IOTYPE_INPUT)
  169. if anode is not None:
  170. output_operand_index = self.add_operand(anode.name, Operand.IOTYPE_OUTPUT)
  171. else:
  172. output_operand_index = self.add_operand(self.edges[bnode.name][0].name, Operand.IOTYPE_OUTPUT)
  173. np.array([input_operand_index, output_operand_index], dtype=np.uint32).tofile(f)
  174. def dump_dense_to_file(self, node, f):
  175. assert(node.op == 'MatMul')
  176. self.layer_number = self.layer_number + 1
  177. self.converted_nodes.add(node.name)
  178. scope_name = TFConverter.get_scope_name(node.name)
  179. #knode for kernel, bnode for bias, anode for activation
  180. knode, bnode, anode = self.get_dense_params(scope_name.split('/')[0])
  181. if bnode is not None:
  182. has_bias = 1
  183. btensor = bnode.attr['value'].tensor
  184. if btensor.tensor_shape.dim[0].size == 1:
  185. bias = struct.pack("f", btensor.float_val[0])
  186. else:
  187. bias = btensor.tensor_content
  188. else:
  189. has_bias = 0
  190. if anode is not None:
  191. activation = anode.op
  192. else:
  193. activation = 'None'
  194. ktensor = knode.attr['value'].tensor
  195. in_channels = ktensor.tensor_shape.dim[0].size
  196. out_channels = ktensor.tensor_shape.dim[1].size
  197. if in_channels * out_channels == 1:
  198. kernel = np.float32(ktensor.float_val[0])
  199. else:
  200. kernel = np.frombuffer(ktensor.tensor_content, dtype=np.float32)
  201. kernel = kernel.reshape(in_channels, out_channels)
  202. kernel = np.transpose(kernel, [1, 0])
  203. np.array([self.op2code[node.op], self.conv_activations[activation], in_channels, out_channels, has_bias], dtype=np.uint32).tofile(f)
  204. kernel.tofile(f)
  205. if has_bias:
  206. f.write(bias)
  207. input_name = self.dense_scopename_inputname_dict[scope_name.split('/')[0]]
  208. input_operand_index = self.add_operand(input_name, Operand.IOTYPE_INPUT)
  209. if anode is not None:
  210. output_operand_index = self.add_operand(anode.name, Operand.IOTYPE_OUTPUT)
  211. else:
  212. if bnode is not None:
  213. output_operand_index = self.add_operand(self.edges[bnode.name][0].name, Operand.IOTYPE_OUTPUT)
  214. else:
  215. output_operand_index = self.add_operand(self.edges[scope_name+'/concat_1'][0].name, Operand.IOTYPE_OUTPUT)
  216. np.array([input_operand_index, output_operand_index], dtype=np.uint32).tofile(f)
  217. def dump_simple_conv2d_to_file(self, node, f):
  218. assert(node.op == 'Conv2D')
  219. self.layer_number = self.layer_number + 1
  220. self.converted_nodes.add(node.name)
  221. node0 = self.name_node_dict[node.input[0]]
  222. node1 = self.name_node_dict[node.input[1]]
  223. if node0.op == 'Const':
  224. knode = node0
  225. input_name = node.input[1]
  226. else:
  227. knode = node1
  228. input_name = node.input[0]
  229. ktensor = knode.attr['value'].tensor
  230. filter_height = ktensor.tensor_shape.dim[0].size
  231. filter_width = ktensor.tensor_shape.dim[1].size
  232. in_channels = ktensor.tensor_shape.dim[2].size
  233. out_channels = ktensor.tensor_shape.dim[3].size
  234. if filter_height * filter_width * in_channels * out_channels == 1:
  235. kernel = np.float32(ktensor.float_val[0])
  236. else:
  237. kernel = np.frombuffer(ktensor.tensor_content, dtype=np.float32)
  238. kernel = kernel.reshape(filter_height, filter_width, in_channels, out_channels)
  239. kernel = np.transpose(kernel, [3, 0, 1, 2])
  240. has_bias = 0
  241. dilation = 1
  242. padding = node.attr['padding'].s.decode("utf-8")
  243. np.array([self.op2code[node.op], dilation, self.conv_paddings[padding], self.conv_activations['None'],
  244. in_channels, out_channels, filter_height, has_bias], dtype=np.uint32).tofile(f)
  245. kernel.tofile(f)
  246. input_operand_index = self.add_operand(input_name, Operand.IOTYPE_INPUT)
  247. output_operand_index = self.add_operand(node.name, Operand.IOTYPE_OUTPUT)
  248. np.array([input_operand_index, output_operand_index], dtype=np.uint32).tofile(f)
  249. def dump_depth2space_to_file(self, node, f):
  250. assert(node.op == 'DepthToSpace')
  251. self.layer_number = self.layer_number + 1
  252. block_size = node.attr['block_size'].i
  253. np.array([self.op2code[node.op], block_size], dtype=np.uint32).tofile(f)
  254. self.converted_nodes.add(node.name)
  255. input_operand_index = self.add_operand(node.input[0], Operand.IOTYPE_INPUT)
  256. output_operand_index = self.add_operand(node.name, Operand.IOTYPE_OUTPUT)
  257. np.array([input_operand_index, output_operand_index], dtype=np.uint32).tofile(f)
  258. def dump_mirrorpad_to_file(self, node, f):
  259. assert(node.op == 'MirrorPad')
  260. self.layer_number = self.layer_number + 1
  261. mode = node.attr['mode'].s
  262. mode = self.mirrorpad_mode[mode.decode("utf-8")]
  263. np.array([self.op2code[node.op], mode], dtype=np.uint32).tofile(f)
  264. pnode = self.name_node_dict[node.input[1]]
  265. self.converted_nodes.add(pnode.name)
  266. paddings = pnode.attr['value'].tensor.tensor_content
  267. f.write(paddings)
  268. self.converted_nodes.add(node.name)
  269. input_operand_index = self.add_operand(node.input[0], Operand.IOTYPE_INPUT)
  270. output_operand_index = self.add_operand(node.name, Operand.IOTYPE_OUTPUT)
  271. np.array([input_operand_index, output_operand_index], dtype=np.uint32).tofile(f)
  272. def dump_maximum_to_file(self, node, f):
  273. assert(node.op == 'Maximum')
  274. self.layer_number = self.layer_number + 1
  275. ynode = self.name_node_dict[node.input[1]]
  276. y = ynode.attr['value'].tensor.float_val[0]
  277. np.array([self.op2code[node.op]], dtype=np.uint32).tofile(f)
  278. np.array([y], dtype=np.float32).tofile(f)
  279. self.converted_nodes.add(node.name)
  280. input_operand_index = self.add_operand(node.input[0], Operand.IOTYPE_INPUT)
  281. output_operand_index = self.add_operand(node.name, Operand.IOTYPE_OUTPUT)
  282. np.array([input_operand_index, output_operand_index], dtype=np.uint32).tofile(f)
  283. def dump_mathbinary_to_file(self, node, f):
  284. self.layer_number = self.layer_number + 1
  285. self.converted_nodes.add(node.name)
  286. i0_node = self.name_node_dict[node.input[0]]
  287. i1_node = self.name_node_dict[node.input[1]]
  288. np.array([self.op2code['MathBinary'], self.mathbin2code[node.op]], dtype=np.uint32).tofile(f)
  289. if i0_node.op == 'Const':
  290. scalar = i0_node.attr['value'].tensor.float_val[0]
  291. np.array([1], dtype=np.uint32).tofile(f) # broadcast: 1
  292. np.array([scalar], dtype=np.float32).tofile(f)
  293. np.array([0], dtype=np.uint32).tofile(f) # broadcast: 0
  294. input_operand_index = self.add_operand(i1_node.name, Operand.IOTYPE_INPUT)
  295. np.array([input_operand_index], dtype=np.uint32).tofile(f)
  296. elif i1_node.op == 'Const':
  297. scalar = i1_node.attr['value'].tensor.float_val[0]
  298. np.array([0], dtype=np.uint32).tofile(f)
  299. input_operand_index = self.add_operand(i0_node.name, Operand.IOTYPE_INPUT)
  300. np.array([input_operand_index], dtype=np.uint32).tofile(f)
  301. np.array([1], dtype=np.uint32).tofile(f)
  302. np.array([scalar], dtype=np.float32).tofile(f)
  303. else:
  304. np.array([0], dtype=np.uint32).tofile(f)
  305. input_operand_index = self.add_operand(i0_node.name, Operand.IOTYPE_INPUT)
  306. np.array([input_operand_index], dtype=np.uint32).tofile(f)
  307. np.array([0], dtype=np.uint32).tofile(f)
  308. input_operand_index = self.add_operand(i1_node.name, Operand.IOTYPE_INPUT)
  309. np.array([input_operand_index], dtype=np.uint32).tofile(f)
  310. output_operand_index = self.add_operand(node.name, Operand.IOTYPE_OUTPUT)
  311. np.array([output_operand_index], dtype=np.uint32).tofile(f)
  312. def dump_mathunary_to_file(self, node, f):
  313. self.layer_number = self.layer_number + 1
  314. self.converted_nodes.add(node.name)
  315. i0_node = self.name_node_dict[node.input[0]]
  316. np.array([self.op2code['MathUnary'], self.mathun2code[node.op]], dtype=np.uint32).tofile(f)
  317. input_operand_index = self.add_operand(i0_node.name, Operand.IOTYPE_INPUT)
  318. np.array([input_operand_index], dtype=np.uint32).tofile(f)
  319. output_operand_index = self.add_operand(node.name, Operand.IOTYPE_OUTPUT)
  320. np.array([output_operand_index],dtype=np.uint32).tofile(f)
  321. def dump_avg_pool_to_file(self, node, f):
  322. assert(node.op == 'AvgPool')
  323. self.layer_number = self.layer_number + 1
  324. self.converted_nodes.add(node.name)
  325. node0 = self.name_node_dict[node.input[0]]
  326. strides = node.attr['strides']
  327. # Tensorflow do not support pooling strides in batch dimension and
  328. # current native NN do not support pooling strides in channel dimension, added assert() here.
  329. assert(strides.list.i[1]==strides.list.i[2])
  330. assert(strides.list.i[0]==1)
  331. assert(strides.list.i[3]==1)
  332. strides = strides.list.i[1]
  333. filter_node = node.attr['ksize']
  334. input_name = node.input[0]
  335. # Tensorflow do not support pooling ksize in batch dimension and channel dimension.
  336. assert(filter_node.list.i[0]==1)
  337. assert(filter_node.list.i[3]==1)
  338. filter_height = filter_node.list.i[1]
  339. filter_width = filter_node.list.i[2]
  340. padding = node.attr['padding'].s.decode("utf-8")
  341. np.array([self.op2code[node.op], strides, self.pool_paddings[padding], filter_height],
  342. dtype=np.uint32).tofile(f)
  343. input_operand_index = self.add_operand(input_name, Operand.IOTYPE_INPUT)
  344. output_operand_index = self.add_operand(node.name, Operand.IOTYPE_OUTPUT)
  345. np.array([input_operand_index, output_operand_index],dtype=np.uint32).tofile(f)
  346. def dump_layers_to_file(self, f):
  347. for node in self.nodes:
  348. if node.name in self.converted_nodes:
  349. continue
  350. # conv2d with dilation generates very complex nodes, so handle it in special
  351. if self.in_conv2d_scope(node.name):
  352. if node.op == 'Conv2D':
  353. self.dump_complex_conv2d_to_file(node, f)
  354. continue
  355. if self.in_dense_scope(node.name):
  356. if node.op == 'MatMul':
  357. self.dump_dense_to_file(node, f)
  358. continue
  359. if node.op == 'Conv2D':
  360. self.dump_simple_conv2d_to_file(node, f)
  361. continue
  362. if node.name in self.output_names:
  363. input_name = self.id_different_scope_dict[node.name]
  364. if TFConverter.get_scope_name(input_name)!=TFConverter.get_scope_name(node.name):
  365. continue
  366. if node.op == 'AvgPool':
  367. self.dump_avg_pool_to_file(node, f)
  368. elif node.op == 'DepthToSpace':
  369. self.dump_depth2space_to_file(node, f)
  370. elif node.op == 'MirrorPad':
  371. self.dump_mirrorpad_to_file(node, f)
  372. elif node.op == 'Maximum':
  373. self.dump_maximum_to_file(node, f)
  374. elif node.op in self.mathbin2code:
  375. self.dump_mathbinary_to_file(node, f)
  376. elif node.op in self.mathun2code:
  377. self.dump_mathunary_to_file(node, f)
  378. def dump_operands_to_file(self, f):
  379. operands = sorted(self.name_operand_dict.values())
  380. for operand in operands:
  381. #print('{}'.format(operand))
  382. np.array([operand.index, len(operand.name)], dtype=np.uint32).tofile(f)
  383. f.write(operand.name.encode('utf-8'))
  384. np.array([operand.iotype, operand.dtype], dtype=np.uint32).tofile(f)
  385. np.array(operand.dims, dtype=np.uint32).tofile(f)
  386. def dump_to_file(self):
  387. with open(self.outfile, 'wb') as f:
  388. f.write(header.str.encode('utf-8'))
  389. np.array([header.major, header.minor], dtype=np.uint32).tofile(f)
  390. self.dump_layers_to_file(f)
  391. self.dump_operands_to_file(f)
  392. np.array([self.layer_number, len(self.name_operand_dict)], dtype=np.uint32).tofile(f)
  393. def generate_name_node_dict(self):
  394. for node in self.nodes:
  395. self.name_node_dict[node.name] = node
  396. def generate_output_names(self):
  397. used_names = []
  398. for node in self.nodes:
  399. for input in node.input:
  400. used_names.append(input)
  401. for node in self.nodes:
  402. if node.name not in used_names:
  403. self.output_names.append(node.name)
  404. def remove_identity(self):
  405. self.id_different_scope_dict = {}
  406. id_nodes = []
  407. id_dict = {}
  408. for node in self.nodes:
  409. if node.op == 'Identity':
  410. name = node.name
  411. input = node.input[0]
  412. id_nodes.append(node)
  413. # do not change the output name
  414. if name in self.output_names:
  415. self.name_node_dict[input].name = name
  416. self.name_node_dict[name] = self.name_node_dict[input]
  417. del self.name_node_dict[input]
  418. self.id_different_scope_dict[name] = input
  419. else:
  420. id_dict[name] = input
  421. for idnode in id_nodes:
  422. self.nodes.remove(idnode)
  423. for node in self.nodes:
  424. for i in range(len(node.input)):
  425. input = node.input[i]
  426. if input in id_dict:
  427. node.input[i] = id_dict[input]
  428. def generate_edges(self):
  429. for node in self.nodes:
  430. for input in node.input:
  431. if input in self.edges:
  432. self.edges[input].append(node)
  433. else:
  434. self.edges[input] = [node]
  435. @staticmethod
  436. def get_scope_name(name):
  437. index = name.rfind('/')
  438. if index == -1:
  439. return ""
  440. return name[0:index]
  441. def in_conv2d_scope(self, name):
  442. inner_scope = TFConverter.get_scope_name(name)
  443. if inner_scope == "":
  444. return False;
  445. for scope in self.conv2d_scope_names:
  446. index = inner_scope.find(scope)
  447. if index == 0:
  448. return True
  449. return False
  450. def in_dense_scope(self, name):
  451. inner_scope = TFConverter.get_scope_name(name)
  452. if inner_scope == "":
  453. return False;
  454. for scope in self.dense_scope_names:
  455. index = inner_scope.find(scope)
  456. if index == 0:
  457. return True
  458. return False
  459. def generate_sub_block_op_scope_info(self):
  460. # mostly, conv2d/dense is a sub block in graph, get the scope name
  461. for node in self.nodes:
  462. if node.op == 'Conv2D':
  463. scope = TFConverter.get_scope_name(node.name)
  464. # for the case tf.nn.conv2d is called directly
  465. if scope == '':
  466. continue
  467. # for the case tf.nn.conv2d is called within a scope
  468. if scope + '/kernel' not in self.name_node_dict:
  469. continue
  470. self.conv2d_scope_names.add(scope)
  471. elif node.op == 'MatMul':
  472. scope = TFConverter.get_scope_name(node.name)
  473. # for the case tf.nn.dense is called directly
  474. if scope == '':
  475. continue
  476. # for the case tf.nn.dense is called within a scope
  477. if scope + '/kernel' not in self.name_node_dict and scope.split('/Tensordot')[0] + '/kernel' not in self.name_node_dict:
  478. continue
  479. self.dense_scope_names.add(scope.split('/Tensordot')[0])
  480. # get the input name to the conv2d/dense sub block
  481. for node in self.nodes:
  482. scope = TFConverter.get_scope_name(node.name)
  483. if scope in self.conv2d_scope_names:
  484. if node.op == 'Conv2D' or node.op == 'Shape':
  485. for inp in node.input:
  486. if TFConverter.get_scope_name(inp) != scope:
  487. self.conv2d_scopename_inputname_dict[scope] = inp
  488. elif scope in self.dense_scope_names:
  489. if node.op == 'MatMul' or node.op == 'Shape':
  490. for inp in node.input:
  491. if TFConverter.get_scope_name(inp) != scope:
  492. self.dense_scopename_inputname_dict[scope] = inp
  493. elif scope.split('/Tensordot')[0] in self.dense_scope_names:
  494. if node.op == 'Transpose':
  495. for inp in node.input:
  496. if TFConverter.get_scope_name(inp).find(scope)<0 and TFConverter.get_scope_name(inp).find(scope.split('/')[0])<0:
  497. self.dense_scopename_inputname_dict[scope.split('/Tensordot')[0]] = inp
  498. def run(self):
  499. self.generate_name_node_dict()
  500. self.generate_output_names()
  501. self.remove_identity()
  502. self.generate_edges()
  503. self.generate_sub_block_op_scope_info()
  504. if self.dump4tb:
  505. self.dump_for_tensorboard()
  506. self.dump_to_file()
  507. def convert_from_tensorflow(infile, outfile, dump4tb):
  508. with open(infile, 'rb') as f:
  509. # read the file in .proto format
  510. graph_def = tf.GraphDef()
  511. graph_def.ParseFromString(f.read())
  512. nodes = graph_def.node
  513. converter = TFConverter(graph_def, nodes, outfile, dump4tb)
  514. converter.run()