Showing the changes to tools/python/convert_from_tensorflow.py (1 of 4 files). Show all
Unified · Split
| @@ -23,6 +23,37 @@ import sys, struct · 22 unmodified lines | ||
| 23 | 23 | |
| 24 | 24 | __all__ = ['convert_from_tensorflow'] |
| 25 | 25 | |
| 26 | +class Operand(object): | |
| 27 | + IOTYPE_INPUT = 1 | |
| 28 | + IOTYPE_OUTPUT = 2 | |
| 29 | + IOTYPE_INTERMEDIATE = IOTYPE_INPUT | IOTYPE_OUTPUT | |
| 30 | + DTYPE_FLOAT = 1 | |
| 31 | + DTYPE_UINT8 = 4 | |
| 32 | + index = 0 | |
| 33 | + def __init__(self, name, dtype, dims): | |
| 34 | + self.name = name | |
| 35 | + self.dtype = dtype | |
| 36 | + self.dims = dims | |
| 37 | + self.iotype = 0 | |
| 38 | + self.used_count = 0 | |
| 39 | + self.index = Operand.index | |
| 40 | + Operand.index = Operand.index + 1 | |
| 41 | + self.iotype2str = {Operand.IOTYPE_INPUT: 'in', Operand.IOTYPE_OUTPUT: 'out', Operand.IOTYPE_INTERMEDIATE: 'inout'} | |
| 42 | + self.dtype2str = {Operand.DTYPE_FLOAT: 'DT_FLOAT', Operand.DTYPE_UINT8: 'DT_UINT8'} | |
| 43 | + | |
| 44 | + def add_iotype(self, iotype): | |
| 45 | + self.iotype = self.iotype | iotype | |
| 46 | + if iotype == Operand.IOTYPE_INPUT: | |
| 47 | + self.used_count = self.used_count + 1 | |
| 48 | + | |
| 49 | + def __str__(self): | |
| 50 | + return "{}: (name: {}, iotype: {}, dtype: {}, dims: ({},{},{},{}) used_count: {})".format(self.index, | |
| 51 | + self.name, self.iotype2str[self.iotype], self.dtype2str[self.dtype], | |
| 52 | + self.dims[0], self.dims[1], self.dims[2], self.dims[3], self.used_count) | |
| 53 | + | |
| 54 | + def __lt__(self, other): | |
| 55 | + return self.index < other.index | |
| 56 | + | |
| 26 | 57 | class TFConverter: |
| 27 | 58 | def __init__(self, graph_def, nodes, outfile, dump4tb): |
| 28 | 59 | self.graph_def = graph_def |
| @@ -37,8 +68,28 @@ class TFConverter: · 8 unmodified lines | ||
| 37 | 68 | self.conv_paddings = {'VALID':0, 'SAME':1} |
| 38 | 69 | self.converted_nodes = set() |
| 39 | 70 | self.conv2d_scope_names = set() |
| 71 | + self.conv2d_scopename_inputname_dict = {} | |
| 40 | 72 | self.op2code = {'Conv2D':1, 'DepthToSpace':2, 'MirrorPad':3} |
| 41 | 73 | self.mirrorpad_mode = {'CONSTANT':0, 'REFLECT':1, 'SYMMETRIC':2} |
| 74 | + self.name_operand_dict = {} | |
| 75 | + | |
| 76 | + | |
| 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 | |
| 42 | 93 | |
| 43 | 94 | |
| 44 | 95 | def dump_for_tensorboard(self): |
| @@ -60,11 +111,10 @@ class TFConverter: · 15 unmodified lines | ||
| 60 | 111 | # the BiasAdd name is possible be changed into the output name, |
| 61 | 112 | # if activation is None, and BiasAdd.next is the last op which is Identity |
| 62 | 113 | if conv2d_scope_name + '/BiasAdd' in self.edges: |
| 63 | − activation = self.edges[conv2d_scope_name + '/BiasAdd'][0] | |
| 64 | − activation = activation.op | |
| 114 | + anode = self.edges[conv2d_scope_name + '/BiasAdd'][0] | |
| 65 | 115 | else: |
| 66 | − activation = 'None' | |
| 67 | − return knode, bnode, dnode, activation | |
| 116 | + anode = None | |
| 117 | + return knode, bnode, dnode, anode | |
| 68 | 118 | |
| 69 | 119 | |
| 70 | 120 | def dump_conv2d_to_file(self, node, f): |
| @@ -73,16 +123,21 @@ class TFConverter: · 2 unmodified lines | ||
| 73 | 123 | self.converted_nodes.add(node.name) |
| 74 | 124 | |
| 75 | 125 | scope_name = TFConverter.get_scope_name(node.name) |
| 76 | − #knode for kernel, bnode for bias, dnode for dilation | |
| 77 | − knode, bnode, dnode, activation = self.get_conv2d_params(scope_name) | |
| 126 | + #knode for kernel, bnode for bias, dnode for dilation, anode for activation | |
| 127 | + knode, bnode, dnode, anode = self.get_conv2d_params(scope_name) | |
| 78 | 128 | |
| 79 | 129 | if dnode is not None: |
| 80 | 130 | dilation = struct.unpack('i', dnode.attr['value'].tensor.tensor_content[0:4])[0] |
| 81 | 131 | else: |
| 82 | 132 | dilation = 1 |
| 83 | 133 | |
| 134 | + if anode is not None: | |
| 135 | + activation = anode.op | |
| 136 | + else: | |
| 137 | + activation = 'None' | |
| 138 | + | |
| 84 | 139 | padding = node.attr['padding'].s.decode("utf-8") |
| 85 | − # conv2d with dilation > 1 generates tens of nodes, not easy to parse them, so use tricky. | |
| 140 | + # conv2d with dilation > 1 generates tens of nodes, not easy to parse them, so use this tricky method. | |
| 86 | 141 | if dilation > 1 and scope_name + '/stack' in self.name_node_dict: |
| 87 | 142 | if self.name_node_dict[scope_name + '/stack'].op == "Const": |
| 88 | 143 | padding = 'SAME' |
| @@ -107,6 +162,15 @@ class TFConverter: · 18 unmodified lines | ||
| 107 | 162 | bias = btensor.tensor_content |
| 108 | 163 | f.write(bias) |
| 109 | 164 | |
| 165 | + input_name = self.conv2d_scopename_inputname_dict[scope_name] | |
| 166 | + input_operand_index = self.add_operand(input_name, Operand.IOTYPE_INPUT) | |
| 167 | + | |
| 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 | + | |
| 110 | 174 | |
| 111 | 175 | def dump_depth2space_to_file(self, node, f): |
| 112 | 176 | assert(node.op == 'DepthToSpace') |
| @@ -114,6 +178,9 @@ class TFConverter: · 1 unmodified line | ||
| 114 | 178 | block_size = node.attr['block_size'].i |
| 115 | 179 | np.array([self.op2code[node.op], block_size], dtype=np.uint32).tofile(f) |
| 116 | 180 | self.converted_nodes.add(node.name) |
| 181 | + input_operand_index = self.add_operand(node.input[0], Operand.IOTYPE_INPUT) | |
| 182 | + output_operand_index = self.add_operand(node.name, Operand.IOTYPE_OUTPUT) | |
| 183 | + np.array([input_operand_index, output_operand_index], dtype=np.uint32).tofile(f) | |
| 117 | 184 | |
| 118 | 185 | |
| 119 | 186 | def dump_mirrorpad_to_file(self, node, f): |
| @@ -127,6 +194,9 @@ class TFConverter: · 7 unmodified lines | ||
| 127 | 194 | paddings = pnode.attr['value'].tensor.tensor_content |
| 128 | 195 | f.write(paddings) |
| 129 | 196 | self.converted_nodes.add(node.name) |
| 197 | + input_operand_index = self.add_operand(node.input[0], Operand.IOTYPE_INPUT) | |
| 198 | + output_operand_index = self.add_operand(node.name, Operand.IOTYPE_OUTPUT) | |
| 199 | + np.array([input_operand_index, output_operand_index], dtype=np.uint32).tofile(f) | |
| 130 | 200 | |
| 131 | 201 | |
| 132 | 202 | def dump_layers_to_file(self, f): |
| @@ -147,10 +217,21 @@ class TFConverter: · 14 unmodified lines | ||
| 147 | 217 | self.dump_mirrorpad_to_file(node, f) |
| 148 | 218 | |
| 149 | 219 | |
| 220 | + def dump_operands_to_file(self, f): | |
| 221 | + operands = sorted(self.name_operand_dict.values()) | |
| 222 | + for operand in operands: | |
| 223 | + #print('{}'.format(operand)) | |
| 224 | + np.array([operand.index, len(operand.name)], dtype=np.uint32).tofile(f) | |
| 225 | + f.write(operand.name.encode('utf-8')) | |
| 226 | + np.array([operand.iotype, operand.dtype], dtype=np.uint32).tofile(f) | |
| 227 | + np.array([operand.dims[0], operand.dims[1], operand.dims[2], operand.dims[3]], dtype=np.uint32).tofile(f) | |
| 228 | + | |
| 229 | + | |
| 150 | 230 | def dump_to_file(self): |
| 151 | 231 | with open(self.outfile, 'wb') as f: |
| 152 | 232 | self.dump_layers_to_file(f) |
| 153 | − np.array([self.layer_number], dtype=np.uint32).tofile(f) | |
| 233 | + self.dump_operands_to_file(f) | |
| 234 | + np.array([self.layer_number, len(self.name_operand_dict)], dtype=np.uint32).tofile(f) | |
| 154 | 235 | |
| 155 | 236 | |
| 156 | 237 | def generate_name_node_dict(self): |
| @@ -212,19 +293,29 @@ class TFConverter: · 55 unmodified lines | ||
| 212 | 293 | return name[0:index] |
| 213 | 294 | |
| 214 | 295 | |
| 215 | − def generate_conv2d_scope_names(self): | |
| 296 | + def generate_conv2d_scope_info(self): | |
| 297 | + # conv2d is a sub block in graph, get the scope name | |
| 216 | 298 | for node in self.nodes: |
| 217 | 299 | if node.op == 'Conv2D': |
| 218 | 300 | scope = TFConverter.get_scope_name(node.name) |
| 219 | 301 | self.conv2d_scope_names.add(scope) |
| 220 | 302 | |
| 303 | + # get the input name to the conv2d sub block | |
| 304 | + for node in self.nodes: | |
| 305 | + scope = TFConverter.get_scope_name(node.name) | |
| 306 | + if scope in self.conv2d_scope_names: | |
| 307 | + if node.op == 'Conv2D' or node.op == 'Shape': | |
| 308 | + for inp in node.input: | |
| 309 | + if TFConverter.get_scope_name(inp) != scope: | |
| 310 | + self.conv2d_scopename_inputname_dict[scope] = inp | |
| 311 | + | |
| 221 | 312 | |
| 222 | 313 | def run(self): |
| 223 | 314 | self.generate_name_node_dict() |
| 224 | 315 | self.generate_output_names() |
| 225 | 316 | self.remove_identity() |
| 226 | 317 | self.generate_edges() |
| 227 | − self.generate_conv2d_scope_names() | |
| 318 | + self.generate_conv2d_scope_info() | |
| 228 | 319 | |
| 229 | 320 | if self.dump4tb: |
| 230 | 321 | self.dump_for_tensorboard() |