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.

335 lines
13KB

  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. __all__ = ['convert_from_tensorflow']
  23. class Operand(object):
  24. IOTYPE_INPUT = 1
  25. IOTYPE_OUTPUT = 2
  26. IOTYPE_INTERMEDIATE = IOTYPE_INPUT | IOTYPE_OUTPUT
  27. DTYPE_FLOAT = 1
  28. DTYPE_UINT8 = 4
  29. index = 0
  30. def __init__(self, name, dtype, dims):
  31. self.name = name
  32. self.dtype = dtype
  33. self.dims = dims
  34. self.iotype = 0
  35. self.used_count = 0
  36. self.index = Operand.index
  37. Operand.index = Operand.index + 1
  38. self.iotype2str = {Operand.IOTYPE_INPUT: 'in', Operand.IOTYPE_OUTPUT: 'out', Operand.IOTYPE_INTERMEDIATE: 'inout'}
  39. self.dtype2str = {Operand.DTYPE_FLOAT: 'DT_FLOAT', Operand.DTYPE_UINT8: 'DT_UINT8'}
  40. def add_iotype(self, iotype):
  41. self.iotype = self.iotype | iotype
  42. if iotype == Operand.IOTYPE_INPUT:
  43. self.used_count = self.used_count + 1
  44. def __str__(self):
  45. return "{}: (name: {}, iotype: {}, dtype: {}, dims: ({},{},{},{}) used_count: {})".format(self.index,
  46. self.name, self.iotype2str[self.iotype], self.dtype2str[self.dtype],
  47. self.dims[0], self.dims[1], self.dims[2], self.dims[3], self.used_count)
  48. def __lt__(self, other):
  49. return self.index < other.index
  50. class TFConverter:
  51. def __init__(self, graph_def, nodes, outfile, dump4tb):
  52. self.graph_def = graph_def
  53. self.nodes = nodes
  54. self.outfile = outfile
  55. self.dump4tb = dump4tb
  56. self.layer_number = 0
  57. self.output_names = []
  58. self.name_node_dict = {}
  59. self.edges = {}
  60. self.conv_activations = {'Relu':0, 'Tanh':1, 'Sigmoid':2, 'None':3, 'LeakyRelu':4}
  61. self.conv_paddings = {'VALID':0, 'SAME':1}
  62. self.converted_nodes = set()
  63. self.conv2d_scope_names = set()
  64. self.conv2d_scopename_inputname_dict = {}
  65. self.op2code = {'Conv2D':1, 'DepthToSpace':2, 'MirrorPad':3}
  66. self.mirrorpad_mode = {'CONSTANT':0, 'REFLECT':1, 'SYMMETRIC':2}
  67. self.name_operand_dict = {}
  68. def add_operand(self, name, type):
  69. node = self.name_node_dict[name]
  70. if name not in self.name_operand_dict:
  71. dtype = node.attr['dtype'].type
  72. if dtype == 0:
  73. dtype = node.attr['T'].type
  74. dims = [-1,-1,-1,-1]
  75. if 'shape' in node.attr:
  76. dims[0] = node.attr['shape'].shape.dim[0].size
  77. dims[1] = node.attr['shape'].shape.dim[1].size
  78. dims[2] = node.attr['shape'].shape.dim[2].size
  79. dims[3] = node.attr['shape'].shape.dim[3].size
  80. operand = Operand(name, dtype, dims)
  81. self.name_operand_dict[name] = operand;
  82. self.name_operand_dict[name].add_iotype(type)
  83. return self.name_operand_dict[name].index
  84. def dump_for_tensorboard(self):
  85. graph = tf.get_default_graph()
  86. tf.import_graph_def(self.graph_def, name="")
  87. tf.summary.FileWriter('/tmp/graph', graph)
  88. print('graph saved, run "tensorboard --logdir=/tmp/graph" to see it')
  89. def get_conv2d_params(self, conv2d_scope_name):
  90. knode = self.name_node_dict[conv2d_scope_name + '/kernel']
  91. bnode = self.name_node_dict[conv2d_scope_name + '/bias']
  92. if conv2d_scope_name + '/dilation_rate' in self.name_node_dict:
  93. dnode = self.name_node_dict[conv2d_scope_name + '/dilation_rate']
  94. else:
  95. dnode = None
  96. # the BiasAdd name is possible be changed into the output name,
  97. # if activation is None, and BiasAdd.next is the last op which is Identity
  98. if conv2d_scope_name + '/BiasAdd' in self.edges:
  99. anode = self.edges[conv2d_scope_name + '/BiasAdd'][0]
  100. else:
  101. anode = None
  102. return knode, bnode, dnode, anode
  103. def dump_conv2d_to_file(self, node, f):
  104. assert(node.op == 'Conv2D')
  105. self.layer_number = self.layer_number + 1
  106. self.converted_nodes.add(node.name)
  107. scope_name = TFConverter.get_scope_name(node.name)
  108. #knode for kernel, bnode for bias, dnode for dilation, anode for activation
  109. knode, bnode, dnode, anode = self.get_conv2d_params(scope_name)
  110. if dnode is not None:
  111. dilation = struct.unpack('i', dnode.attr['value'].tensor.tensor_content[0:4])[0]
  112. else:
  113. dilation = 1
  114. if anode is not None:
  115. activation = anode.op
  116. else:
  117. activation = 'None'
  118. padding = node.attr['padding'].s.decode("utf-8")
  119. # conv2d with dilation > 1 generates tens of nodes, not easy to parse them, so use this tricky method.
  120. if dilation > 1 and scope_name + '/stack' in self.name_node_dict:
  121. if self.name_node_dict[scope_name + '/stack'].op == "Const":
  122. padding = 'SAME'
  123. padding = self.conv_paddings[padding]
  124. ktensor = knode.attr['value'].tensor
  125. filter_height = ktensor.tensor_shape.dim[0].size
  126. filter_width = ktensor.tensor_shape.dim[1].size
  127. in_channels = ktensor.tensor_shape.dim[2].size
  128. out_channels = ktensor.tensor_shape.dim[3].size
  129. kernel = np.frombuffer(ktensor.tensor_content, dtype=np.float32)
  130. kernel = kernel.reshape(filter_height, filter_width, in_channels, out_channels)
  131. kernel = np.transpose(kernel, [3, 0, 1, 2])
  132. np.array([self.op2code[node.op], dilation, padding, self.conv_activations[activation], in_channels, out_channels, filter_height], dtype=np.uint32).tofile(f)
  133. kernel.tofile(f)
  134. btensor = bnode.attr['value'].tensor
  135. if btensor.tensor_shape.dim[0].size == 1:
  136. bias = struct.pack("f", btensor.float_val[0])
  137. else:
  138. bias = btensor.tensor_content
  139. f.write(bias)
  140. input_name = self.conv2d_scopename_inputname_dict[scope_name]
  141. input_operand_index = self.add_operand(input_name, Operand.IOTYPE_INPUT)
  142. if anode is not None:
  143. output_operand_index = self.add_operand(anode.name, Operand.IOTYPE_OUTPUT)
  144. else:
  145. output_operand_index = self.add_operand(self.edges[bnode.name][0].name, Operand.IOTYPE_OUTPUT)
  146. np.array([input_operand_index, output_operand_index], dtype=np.uint32).tofile(f)
  147. def dump_depth2space_to_file(self, node, f):
  148. assert(node.op == 'DepthToSpace')
  149. self.layer_number = self.layer_number + 1
  150. block_size = node.attr['block_size'].i
  151. np.array([self.op2code[node.op], block_size], dtype=np.uint32).tofile(f)
  152. self.converted_nodes.add(node.name)
  153. input_operand_index = self.add_operand(node.input[0], Operand.IOTYPE_INPUT)
  154. output_operand_index = self.add_operand(node.name, Operand.IOTYPE_OUTPUT)
  155. np.array([input_operand_index, output_operand_index], dtype=np.uint32).tofile(f)
  156. def dump_mirrorpad_to_file(self, node, f):
  157. assert(node.op == 'MirrorPad')
  158. self.layer_number = self.layer_number + 1
  159. mode = node.attr['mode'].s
  160. mode = self.mirrorpad_mode[mode.decode("utf-8")]
  161. np.array([self.op2code[node.op], mode], dtype=np.uint32).tofile(f)
  162. pnode = self.name_node_dict[node.input[1]]
  163. self.converted_nodes.add(pnode.name)
  164. paddings = pnode.attr['value'].tensor.tensor_content
  165. f.write(paddings)
  166. self.converted_nodes.add(node.name)
  167. input_operand_index = self.add_operand(node.input[0], Operand.IOTYPE_INPUT)
  168. output_operand_index = self.add_operand(node.name, Operand.IOTYPE_OUTPUT)
  169. np.array([input_operand_index, output_operand_index], dtype=np.uint32).tofile(f)
  170. def dump_layers_to_file(self, f):
  171. for node in self.nodes:
  172. if node.name in self.converted_nodes:
  173. continue
  174. # conv2d with dilation generates very complex nodes, so handle it in special
  175. scope_name = TFConverter.get_scope_name(node.name)
  176. if scope_name in self.conv2d_scope_names:
  177. if node.op == 'Conv2D':
  178. self.dump_conv2d_to_file(node, f)
  179. continue
  180. if node.op == 'DepthToSpace':
  181. self.dump_depth2space_to_file(node, f)
  182. elif node.op == 'MirrorPad':
  183. self.dump_mirrorpad_to_file(node, f)
  184. def dump_operands_to_file(self, f):
  185. operands = sorted(self.name_operand_dict.values())
  186. for operand in operands:
  187. #print('{}'.format(operand))
  188. np.array([operand.index, len(operand.name)], dtype=np.uint32).tofile(f)
  189. f.write(operand.name.encode('utf-8'))
  190. np.array([operand.iotype, operand.dtype], dtype=np.uint32).tofile(f)
  191. np.array([operand.dims[0], operand.dims[1], operand.dims[2], operand.dims[3]], dtype=np.uint32).tofile(f)
  192. def dump_to_file(self):
  193. with open(self.outfile, 'wb') as f:
  194. self.dump_layers_to_file(f)
  195. self.dump_operands_to_file(f)
  196. np.array([self.layer_number, len(self.name_operand_dict)], dtype=np.uint32).tofile(f)
  197. def generate_name_node_dict(self):
  198. for node in self.nodes:
  199. self.name_node_dict[node.name] = node
  200. def generate_output_names(self):
  201. used_names = []
  202. for node in self.nodes:
  203. for input in node.input:
  204. used_names.append(input)
  205. for node in self.nodes:
  206. if node.name not in used_names:
  207. self.output_names.append(node.name)
  208. def remove_identity(self):
  209. id_nodes = []
  210. id_dict = {}
  211. for node in self.nodes:
  212. if node.op == 'Identity':
  213. name = node.name
  214. input = node.input[0]
  215. id_nodes.append(node)
  216. # do not change the output name
  217. if name in self.output_names:
  218. self.name_node_dict[input].name = name
  219. self.name_node_dict[name] = self.name_node_dict[input]
  220. del self.name_node_dict[input]
  221. else:
  222. id_dict[name] = input
  223. for idnode in id_nodes:
  224. self.nodes.remove(idnode)
  225. for node in self.nodes:
  226. for i in range(len(node.input)):
  227. input = node.input[i]
  228. if input in id_dict:
  229. node.input[i] = id_dict[input]
  230. def generate_edges(self):
  231. for node in self.nodes:
  232. for input in node.input:
  233. if input in self.edges:
  234. self.edges[input].append(node)
  235. else:
  236. self.edges[input] = [node]
  237. @staticmethod
  238. def get_scope_name(name):
  239. index = name.rfind('/')
  240. if index == -1:
  241. return ""
  242. return name[0:index]
  243. def generate_conv2d_scope_info(self):
  244. # conv2d is a sub block in graph, get the scope name
  245. for node in self.nodes:
  246. if node.op == 'Conv2D':
  247. scope = TFConverter.get_scope_name(node.name)
  248. self.conv2d_scope_names.add(scope)
  249. # get the input name to the conv2d sub block
  250. for node in self.nodes:
  251. scope = TFConverter.get_scope_name(node.name)
  252. if scope in self.conv2d_scope_names:
  253. if node.op == 'Conv2D' or node.op == 'Shape':
  254. for inp in node.input:
  255. if TFConverter.get_scope_name(inp) != scope:
  256. self.conv2d_scopename_inputname_dict[scope] = inp
  257. def run(self):
  258. self.generate_name_node_dict()
  259. self.generate_output_names()
  260. self.remove_identity()
  261. self.generate_edges()
  262. self.generate_conv2d_scope_info()
  263. if self.dump4tb:
  264. self.dump_for_tensorboard()
  265. self.dump_to_file()
  266. def convert_from_tensorflow(infile, outfile, dump4tb):
  267. with open(infile, 'rb') as f:
  268. # read the file in .proto format
  269. graph_def = tf.GraphDef()
  270. graph_def.ParseFromString(f.read())
  271. nodes = graph_def.node
  272. converter = TFConverter(graph_def, nodes, outfile, dump4tb)
  273. converter.run()