You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

607 lines
25KB

  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. self.mirrorpad_mode = {'CONSTANT':0, 'REFLECT':1, 'SYMMETRIC':2}
  76. self.name_operand_dict = {}
  77. def add_operand(self, name, type):
  78. node = self.name_node_dict[name]
  79. if name not in self.name_operand_dict:
  80. dtype = node.attr['dtype'].type
  81. if dtype == 0:
  82. dtype = node.attr['T'].type
  83. dims = [-1,-1,-1,-1]
  84. if 'shape' in node.attr:
  85. dims[0] = node.attr['shape'].shape.dim[0].size
  86. dims[1] = node.attr['shape'].shape.dim[1].size
  87. dims[2] = node.attr['shape'].shape.dim[2].size
  88. dims[3] = node.attr['shape'].shape.dim[3].size
  89. operand = Operand(name, dtype, dims)
  90. self.name_operand_dict[name] = operand;
  91. self.name_operand_dict[name].add_iotype(type)
  92. return self.name_operand_dict[name].index
  93. def dump_for_tensorboard(self):
  94. graph = tf.get_default_graph()
  95. tf.import_graph_def(self.graph_def, name="")
  96. tf.summary.FileWriter('/tmp/graph', graph)
  97. print('graph saved, run "tensorboard --logdir=/tmp/graph" to see it')
  98. def get_conv2d_params(self, conv2d_scope_name):
  99. knode = self.name_node_dict[conv2d_scope_name + '/kernel']
  100. bnode = self.name_node_dict[conv2d_scope_name + '/bias']
  101. if conv2d_scope_name + '/dilation_rate' in self.name_node_dict:
  102. dnode = self.name_node_dict[conv2d_scope_name + '/dilation_rate']
  103. else:
  104. dnode = None
  105. # the BiasAdd name is possible be changed into the output name,
  106. # if activation is None, and BiasAdd.next is the last op which is Identity
  107. if conv2d_scope_name + '/BiasAdd' in self.edges:
  108. anode = self.edges[conv2d_scope_name + '/BiasAdd'][0]
  109. if anode.op not in self.conv_activations:
  110. anode = None
  111. else:
  112. anode = None
  113. return knode, bnode, dnode, anode
  114. def get_dense_params(self, dense_scope_name):
  115. knode = self.name_node_dict[dense_scope_name + '/kernel']
  116. bnode = self.name_node_dict.get(dense_scope_name + '/bias')
  117. # the BiasAdd name is possible be changed into the output name,
  118. # if activation is None, and BiasAdd.next is the last op which is Identity
  119. anode = None
  120. if bnode:
  121. if dense_scope_name + '/BiasAdd' in self.edges:
  122. anode = self.edges[dense_scope_name + '/BiasAdd'][0]
  123. if anode.op not in self.conv_activations:
  124. anode = None
  125. else:
  126. anode = None
  127. return knode, bnode, anode
  128. def dump_complex_conv2d_to_file(self, node, f):
  129. assert(node.op == 'Conv2D')
  130. self.layer_number = self.layer_number + 1
  131. self.converted_nodes.add(node.name)
  132. scope_name = TFConverter.get_scope_name(node.name)
  133. #knode for kernel, bnode for bias, dnode for dilation, anode for activation
  134. knode, bnode, dnode, anode = self.get_conv2d_params(scope_name)
  135. if dnode is not None:
  136. dilation = struct.unpack('i', dnode.attr['value'].tensor.tensor_content[0:4])[0]
  137. else:
  138. dilation = 1
  139. if anode is not None:
  140. activation = anode.op
  141. else:
  142. activation = 'None'
  143. padding = node.attr['padding'].s.decode("utf-8")
  144. # conv2d with dilation > 1 generates tens of nodes, not easy to parse them, so use this tricky method.
  145. if dilation > 1 and scope_name + '/stack' in self.name_node_dict:
  146. if self.name_node_dict[scope_name + '/stack'].op == "Const":
  147. padding = 'SAME'
  148. padding = self.conv_paddings[padding]
  149. ktensor = knode.attr['value'].tensor
  150. filter_height = ktensor.tensor_shape.dim[0].size
  151. filter_width = ktensor.tensor_shape.dim[1].size
  152. in_channels = ktensor.tensor_shape.dim[2].size
  153. out_channels = ktensor.tensor_shape.dim[3].size
  154. kernel = np.frombuffer(ktensor.tensor_content, dtype=np.float32)
  155. kernel = kernel.reshape(filter_height, filter_width, in_channels, out_channels)
  156. kernel = np.transpose(kernel, [3, 0, 1, 2])
  157. has_bias = 1
  158. 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)
  159. kernel.tofile(f)
  160. btensor = bnode.attr['value'].tensor
  161. if btensor.tensor_shape.dim[0].size == 1:
  162. bias = struct.pack("f", btensor.float_val[0])
  163. else:
  164. bias = btensor.tensor_content
  165. f.write(bias)
  166. input_name = self.conv2d_scopename_inputname_dict[scope_name]
  167. input_operand_index = self.add_operand(input_name, Operand.IOTYPE_INPUT)
  168. if anode is not None:
  169. output_operand_index = self.add_operand(anode.name, Operand.IOTYPE_OUTPUT)
  170. else:
  171. output_operand_index = self.add_operand(self.edges[bnode.name][0].name, Operand.IOTYPE_OUTPUT)
  172. np.array([input_operand_index, output_operand_index], dtype=np.uint32).tofile(f)
  173. def dump_dense_to_file(self, node, f):
  174. assert(node.op == 'MatMul')
  175. self.layer_number = self.layer_number + 1
  176. self.converted_nodes.add(node.name)
  177. scope_name = TFConverter.get_scope_name(node.name)
  178. #knode for kernel, bnode for bias, anode for activation
  179. knode, bnode, anode = self.get_dense_params(scope_name.split('/')[0])
  180. if bnode is not None:
  181. has_bias = 1
  182. btensor = bnode.attr['value'].tensor
  183. if btensor.tensor_shape.dim[0].size == 1:
  184. bias = struct.pack("f", btensor.float_val[0])
  185. else:
  186. bias = btensor.tensor_content
  187. else:
  188. has_bias = 0
  189. if anode is not None:
  190. activation = anode.op
  191. else:
  192. activation = 'None'
  193. ktensor = knode.attr['value'].tensor
  194. in_channels = ktensor.tensor_shape.dim[0].size
  195. out_channels = ktensor.tensor_shape.dim[1].size
  196. if in_channels * out_channels == 1:
  197. kernel = np.float32(ktensor.float_val[0])
  198. else:
  199. kernel = np.frombuffer(ktensor.tensor_content, dtype=np.float32)
  200. kernel = kernel.reshape(in_channels, out_channels)
  201. kernel = np.transpose(kernel, [1, 0])
  202. np.array([self.op2code[node.op], self.conv_activations[activation], in_channels, out_channels, has_bias], dtype=np.uint32).tofile(f)
  203. kernel.tofile(f)
  204. if has_bias:
  205. f.write(bias)
  206. input_name = self.dense_scopename_inputname_dict[scope_name.split('/')[0]]
  207. input_operand_index = self.add_operand(input_name, Operand.IOTYPE_INPUT)
  208. if anode is not None:
  209. output_operand_index = self.add_operand(anode.name, Operand.IOTYPE_OUTPUT)
  210. else:
  211. if bnode is not None:
  212. output_operand_index = self.add_operand(self.edges[bnode.name][0].name, Operand.IOTYPE_OUTPUT)
  213. else:
  214. output_operand_index = self.add_operand(self.edges[scope_name+'/concat_1'][0].name, Operand.IOTYPE_OUTPUT)
  215. np.array([input_operand_index, output_operand_index], dtype=np.uint32).tofile(f)
  216. def dump_simple_conv2d_to_file(self, node, f):
  217. assert(node.op == 'Conv2D')
  218. self.layer_number = self.layer_number + 1
  219. self.converted_nodes.add(node.name)
  220. node0 = self.name_node_dict[node.input[0]]
  221. node1 = self.name_node_dict[node.input[1]]
  222. if node0.op == 'Const':
  223. knode = node0
  224. input_name = node.input[1]
  225. else:
  226. knode = node1
  227. input_name = node.input[0]
  228. ktensor = knode.attr['value'].tensor
  229. filter_height = ktensor.tensor_shape.dim[0].size
  230. filter_width = ktensor.tensor_shape.dim[1].size
  231. in_channels = ktensor.tensor_shape.dim[2].size
  232. out_channels = ktensor.tensor_shape.dim[3].size
  233. if filter_height * filter_width * in_channels * out_channels == 1:
  234. kernel = np.float32(ktensor.float_val[0])
  235. else:
  236. kernel = np.frombuffer(ktensor.tensor_content, dtype=np.float32)
  237. kernel = kernel.reshape(filter_height, filter_width, in_channels, out_channels)
  238. kernel = np.transpose(kernel, [3, 0, 1, 2])
  239. has_bias = 0
  240. dilation = 1
  241. padding = node.attr['padding'].s.decode("utf-8")
  242. np.array([self.op2code[node.op], dilation, self.conv_paddings[padding], self.conv_activations['None'],
  243. in_channels, out_channels, filter_height, has_bias], dtype=np.uint32).tofile(f)
  244. kernel.tofile(f)
  245. input_operand_index = self.add_operand(input_name, Operand.IOTYPE_INPUT)
  246. output_operand_index = self.add_operand(node.name, Operand.IOTYPE_OUTPUT)
  247. np.array([input_operand_index, output_operand_index], dtype=np.uint32).tofile(f)
  248. def dump_depth2space_to_file(self, node, f):
  249. assert(node.op == 'DepthToSpace')
  250. self.layer_number = self.layer_number + 1
  251. block_size = node.attr['block_size'].i
  252. np.array([self.op2code[node.op], block_size], dtype=np.uint32).tofile(f)
  253. self.converted_nodes.add(node.name)
  254. input_operand_index = self.add_operand(node.input[0], Operand.IOTYPE_INPUT)
  255. output_operand_index = self.add_operand(node.name, Operand.IOTYPE_OUTPUT)
  256. np.array([input_operand_index, output_operand_index], dtype=np.uint32).tofile(f)
  257. def dump_mirrorpad_to_file(self, node, f):
  258. assert(node.op == 'MirrorPad')
  259. self.layer_number = self.layer_number + 1
  260. mode = node.attr['mode'].s
  261. mode = self.mirrorpad_mode[mode.decode("utf-8")]
  262. np.array([self.op2code[node.op], mode], dtype=np.uint32).tofile(f)
  263. pnode = self.name_node_dict[node.input[1]]
  264. self.converted_nodes.add(pnode.name)
  265. paddings = pnode.attr['value'].tensor.tensor_content
  266. f.write(paddings)
  267. self.converted_nodes.add(node.name)
  268. input_operand_index = self.add_operand(node.input[0], Operand.IOTYPE_INPUT)
  269. output_operand_index = self.add_operand(node.name, Operand.IOTYPE_OUTPUT)
  270. np.array([input_operand_index, output_operand_index], dtype=np.uint32).tofile(f)
  271. def dump_maximum_to_file(self, node, f):
  272. assert(node.op == 'Maximum')
  273. self.layer_number = self.layer_number + 1
  274. ynode = self.name_node_dict[node.input[1]]
  275. y = ynode.attr['value'].tensor.float_val[0]
  276. np.array([self.op2code[node.op]], dtype=np.uint32).tofile(f)
  277. np.array([y], dtype=np.float32).tofile(f)
  278. self.converted_nodes.add(node.name)
  279. input_operand_index = self.add_operand(node.input[0], Operand.IOTYPE_INPUT)
  280. output_operand_index = self.add_operand(node.name, Operand.IOTYPE_OUTPUT)
  281. np.array([input_operand_index, output_operand_index], dtype=np.uint32).tofile(f)
  282. def dump_mathbinary_to_file(self, node, f):
  283. self.layer_number = self.layer_number + 1
  284. self.converted_nodes.add(node.name)
  285. i0_node = self.name_node_dict[node.input[0]]
  286. i1_node = self.name_node_dict[node.input[1]]
  287. np.array([self.op2code['MathBinary'], self.mathbin2code[node.op]], dtype=np.uint32).tofile(f)
  288. if i0_node.op == 'Const':
  289. scalar = i0_node.attr['value'].tensor.float_val[0]
  290. np.array([1], dtype=np.uint32).tofile(f) # broadcast: 1
  291. np.array([scalar], dtype=np.float32).tofile(f)
  292. np.array([0], dtype=np.uint32).tofile(f) # broadcast: 0
  293. input_operand_index = self.add_operand(i1_node.name, Operand.IOTYPE_INPUT)
  294. np.array([input_operand_index], dtype=np.uint32).tofile(f)
  295. elif i1_node.op == 'Const':
  296. scalar = i1_node.attr['value'].tensor.float_val[0]
  297. np.array([0], dtype=np.uint32).tofile(f)
  298. input_operand_index = self.add_operand(i0_node.name, Operand.IOTYPE_INPUT)
  299. np.array([input_operand_index], dtype=np.uint32).tofile(f)
  300. np.array([1], dtype=np.uint32).tofile(f)
  301. np.array([scalar], dtype=np.float32).tofile(f)
  302. else:
  303. np.array([0], dtype=np.uint32).tofile(f)
  304. input_operand_index = self.add_operand(i0_node.name, Operand.IOTYPE_INPUT)
  305. np.array([input_operand_index], dtype=np.uint32).tofile(f)
  306. np.array([0], dtype=np.uint32).tofile(f)
  307. input_operand_index = self.add_operand(i1_node.name, Operand.IOTYPE_INPUT)
  308. np.array([input_operand_index], dtype=np.uint32).tofile(f)
  309. output_operand_index = self.add_operand(node.name, Operand.IOTYPE_OUTPUT)
  310. np.array([output_operand_index], dtype=np.uint32).tofile(f)
  311. def dump_mathunary_to_file(self, node, f):
  312. self.layer_number = self.layer_number + 1
  313. self.converted_nodes.add(node.name)
  314. i0_node = self.name_node_dict[node.input[0]]
  315. np.array([self.op2code['MathUnary'], self.mathun2code[node.op]], dtype=np.uint32).tofile(f)
  316. input_operand_index = self.add_operand(i0_node.name, Operand.IOTYPE_INPUT)
  317. np.array([input_operand_index], dtype=np.uint32).tofile(f)
  318. output_operand_index = self.add_operand(node.name, Operand.IOTYPE_OUTPUT)
  319. np.array([output_operand_index],dtype=np.uint32).tofile(f)
  320. def dump_avg_pool_to_file(self, node, f):
  321. assert(node.op == 'AvgPool')
  322. self.layer_number = self.layer_number + 1
  323. self.converted_nodes.add(node.name)
  324. node0 = self.name_node_dict[node.input[0]]
  325. strides = node.attr['strides']
  326. # Tensorflow do not support pooling strides in batch dimension and
  327. # current native NN do not support pooling strides in channel dimension, added assert() here.
  328. assert(strides.list.i[1]==strides.list.i[2])
  329. assert(strides.list.i[0]==1)
  330. assert(strides.list.i[3]==1)
  331. strides = strides.list.i[1]
  332. filter_node = node.attr['ksize']
  333. input_name = node.input[0]
  334. # Tensorflow do not support pooling ksize in batch dimension and channel dimension.
  335. assert(filter_node.list.i[0]==1)
  336. assert(filter_node.list.i[3]==1)
  337. filter_height = filter_node.list.i[1]
  338. filter_width = filter_node.list.i[2]
  339. padding = node.attr['padding'].s.decode("utf-8")
  340. np.array([self.op2code[node.op], strides, self.pool_paddings[padding], filter_height],
  341. dtype=np.uint32).tofile(f)
  342. input_operand_index = self.add_operand(input_name, Operand.IOTYPE_INPUT)
  343. output_operand_index = self.add_operand(node.name, Operand.IOTYPE_OUTPUT)
  344. np.array([input_operand_index, output_operand_index],dtype=np.uint32).tofile(f)
  345. def dump_layers_to_file(self, f):
  346. for node in self.nodes:
  347. if node.name in self.converted_nodes:
  348. continue
  349. # conv2d with dilation generates very complex nodes, so handle it in special
  350. if self.in_conv2d_scope(node.name):
  351. if node.op == 'Conv2D':
  352. self.dump_complex_conv2d_to_file(node, f)
  353. continue
  354. if self.in_dense_scope(node.name):
  355. if node.op == 'MatMul':
  356. self.dump_dense_to_file(node, f)
  357. continue
  358. if node.op == 'Conv2D':
  359. self.dump_simple_conv2d_to_file(node, f)
  360. continue
  361. if node.name in self.output_names:
  362. input_name = self.id_different_scope_dict[node.name]
  363. if TFConverter.get_scope_name(input_name)!=TFConverter.get_scope_name(node.name):
  364. continue
  365. if node.op == 'AvgPool':
  366. self.dump_avg_pool_to_file(node, f)
  367. elif node.op == 'DepthToSpace':
  368. self.dump_depth2space_to_file(node, f)
  369. elif node.op == 'MirrorPad':
  370. self.dump_mirrorpad_to_file(node, f)
  371. elif node.op == 'Maximum':
  372. self.dump_maximum_to_file(node, f)
  373. elif node.op in self.mathbin2code:
  374. self.dump_mathbinary_to_file(node, f)
  375. elif node.op in self.mathun2code:
  376. self.dump_mathunary_to_file(node, f)
  377. def dump_operands_to_file(self, f):
  378. operands = sorted(self.name_operand_dict.values())
  379. for operand in operands:
  380. #print('{}'.format(operand))
  381. np.array([operand.index, len(operand.name)], dtype=np.uint32).tofile(f)
  382. f.write(operand.name.encode('utf-8'))
  383. np.array([operand.iotype, operand.dtype], dtype=np.uint32).tofile(f)
  384. np.array(operand.dims, dtype=np.uint32).tofile(f)
  385. def dump_to_file(self):
  386. with open(self.outfile, 'wb') as f:
  387. f.write(header.str.encode('utf-8'))
  388. np.array([header.major, header.minor], dtype=np.uint32).tofile(f)
  389. self.dump_layers_to_file(f)
  390. self.dump_operands_to_file(f)
  391. np.array([self.layer_number, len(self.name_operand_dict)], dtype=np.uint32).tofile(f)
  392. def generate_name_node_dict(self):
  393. for node in self.nodes:
  394. self.name_node_dict[node.name] = node
  395. def generate_output_names(self):
  396. used_names = []
  397. for node in self.nodes:
  398. for input in node.input:
  399. used_names.append(input)
  400. for node in self.nodes:
  401. if node.name not in used_names:
  402. self.output_names.append(node.name)
  403. def remove_identity(self):
  404. self.id_different_scope_dict = {}
  405. id_nodes = []
  406. id_dict = {}
  407. for node in self.nodes:
  408. if node.op == 'Identity':
  409. name = node.name
  410. input = node.input[0]
  411. id_nodes.append(node)
  412. # do not change the output name
  413. if name in self.output_names:
  414. self.name_node_dict[input].name = name
  415. self.name_node_dict[name] = self.name_node_dict[input]
  416. del self.name_node_dict[input]
  417. self.id_different_scope_dict[name] = input
  418. else:
  419. id_dict[name] = input
  420. for idnode in id_nodes:
  421. self.nodes.remove(idnode)
  422. for node in self.nodes:
  423. for i in range(len(node.input)):
  424. input = node.input[i]
  425. if input in id_dict:
  426. node.input[i] = id_dict[input]
  427. def generate_edges(self):
  428. for node in self.nodes:
  429. for input in node.input:
  430. if input in self.edges:
  431. self.edges[input].append(node)
  432. else:
  433. self.edges[input] = [node]
  434. @staticmethod
  435. def get_scope_name(name):
  436. index = name.rfind('/')
  437. if index == -1:
  438. return ""
  439. return name[0:index]
  440. def in_conv2d_scope(self, name):
  441. inner_scope = TFConverter.get_scope_name(name)
  442. if inner_scope == "":
  443. return False;
  444. for scope in self.conv2d_scope_names:
  445. index = inner_scope.find(scope)
  446. if index == 0:
  447. return True
  448. return False
  449. def in_dense_scope(self, name):
  450. inner_scope = TFConverter.get_scope_name(name)
  451. if inner_scope == "":
  452. return False;
  453. for scope in self.dense_scope_names:
  454. index = inner_scope.find(scope)
  455. if index == 0:
  456. return True
  457. return False
  458. def generate_sub_block_op_scope_info(self):
  459. # mostly, conv2d/dense is a sub block in graph, get the scope name
  460. for node in self.nodes:
  461. if node.op == 'Conv2D':
  462. scope = TFConverter.get_scope_name(node.name)
  463. # for the case tf.nn.conv2d is called directly
  464. if scope == '':
  465. continue
  466. # for the case tf.nn.conv2d is called within a scope
  467. if scope + '/kernel' not in self.name_node_dict:
  468. continue
  469. self.conv2d_scope_names.add(scope)
  470. elif node.op == 'MatMul':
  471. scope = TFConverter.get_scope_name(node.name)
  472. # for the case tf.nn.dense is called directly
  473. if scope == '':
  474. continue
  475. # for the case tf.nn.dense is called within a scope
  476. if scope + '/kernel' not in self.name_node_dict and scope.split('/Tensordot')[0] + '/kernel' not in self.name_node_dict:
  477. continue
  478. self.dense_scope_names.add(scope.split('/Tensordot')[0])
  479. # get the input name to the conv2d/dense sub block
  480. for node in self.nodes:
  481. scope = TFConverter.get_scope_name(node.name)
  482. if scope in self.conv2d_scope_names:
  483. if node.op == 'Conv2D' or node.op == 'Shape':
  484. for inp in node.input:
  485. if TFConverter.get_scope_name(inp) != scope:
  486. self.conv2d_scopename_inputname_dict[scope] = inp
  487. elif scope in self.dense_scope_names:
  488. if node.op == 'MatMul' or node.op == 'Shape':
  489. for inp in node.input:
  490. if TFConverter.get_scope_name(inp) != scope:
  491. self.dense_scopename_inputname_dict[scope] = inp
  492. elif scope.split('/Tensordot')[0] in self.dense_scope_names:
  493. if node.op == 'Transpose':
  494. for inp in node.input:
  495. if TFConverter.get_scope_name(inp).find(scope)<0 and TFConverter.get_scope_name(inp).find(scope.split('/')[0])<0:
  496. self.dense_scopename_inputname_dict[scope.split('/Tensordot')[0]] = inp
  497. def run(self):
  498. self.generate_name_node_dict()
  499. self.generate_output_names()
  500. self.remove_identity()
  501. self.generate_edges()
  502. self.generate_sub_block_op_scope_info()
  503. if self.dump4tb:
  504. self.dump_for_tensorboard()
  505. self.dump_to_file()
  506. def convert_from_tensorflow(infile, outfile, dump4tb):
  507. with open(infile, 'rb') as f:
  508. # read the file in .proto format
  509. graph_def = tf.GraphDef()
  510. graph_def.ParseFromString(f.read())
  511. nodes = graph_def.node
  512. converter = TFConverter(graph_def, nodes, outfile, dump4tb)
  513. converter.run()