Showing the changes to tools/python (1 of 4 files). Show all
Unified · Split
| @@ -23,6 +23,37 @@ import sys, struct · 22 unmodified lines | |||
| 23 | 23 | ||
| 24 | __all__ = ['convert_from_tensorflow'] | 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 | class TFConverter: | 57 | class TFConverter: |
| 27 | def __init__(self, graph_def, nodes, outfile, dump4tb): | 58 | def __init__(self, graph_def, nodes, outfile, dump4tb): |
| 28 | self.graph_def = graph_def | 59 | self.graph_def = graph_def |
| @@ -37,8 +68,28 @@ class TFConverter: · 8 unmodified lines | |||
| 37 | self.conv_paddings = {'VALID':0, 'SAME':1} | 68 | self.conv_paddings = {'VALID':0, 'SAME':1} |
| 38 | self.converted_nodes = set() | 69 | self.converted_nodes = set() |
| 39 | self.conv2d_scope_names = set() | 70 | self.conv2d_scope_names = set() |
| 71 | self.conv2d_scopename_inputname_dict = {} | ||
| 40 | self.op2code = {'Conv2D':1, 'DepthToSpace':2, 'MirrorPad':3} | 72 | self.op2code = {'Conv2D':1, 'DepthToSpace':2, 'MirrorPad':3} |
| 41 | self.mirrorpad_mode = {'CONSTANT':0, 'REFLECT':1, 'SYMMETRIC':2} | 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 | def dump_for_tensorboard(self): | 95 | def dump_for_tensorboard(self): |
| @@ -60,11 +111,10 @@ class TFConverter: · 15 unmodified lines | |||
| 60 | # the BiasAdd name is possible be changed into the output name, | 111 | # the BiasAdd name is possible be changed into the output name, |
| 61 | # if activation is None, and BiasAdd.next is the last op which is Identity | 112 | # if activation is None, and BiasAdd.next is the last op which is Identity |
| 62 | if conv2d_scope_name + '/BiasAdd' in self.edges: | 113 | if conv2d_scope_name + '/BiasAdd' in self.edges: |
| 63 | activation = self.edges[conv2d_scope_name + '/BiasAdd'][0] | 114 | anode = self.edges[conv2d_scope_name + '/BiasAdd'][0] |
| 64 | activation = activation.op | ||
| 65 | else: | 115 | else: |
| 66 | activation = 'None' | 116 | anode = None |
| 67 | return knode, bnode, dnode, activation | 117 | return knode, bnode, dnode, anode |
| 68 | 118 | ||
| 69 | 119 | ||
| 70 | def dump_conv2d_to_file(self, node, f): | 120 | def dump_conv2d_to_file(self, node, f): |
| @@ -73,16 +123,21 @@ class TFConverter: · 2 unmodified lines | |||
| 73 | self.converted_nodes.add(node.name) | 123 | self.converted_nodes.add(node.name) |
| 74 | 124 | ||
| 75 | scope_name = TFConverter.get_scope_name(node.name) | 125 | scope_name = TFConverter.get_scope_name(node.name) |
| 76 | #knode for kernel, bnode for bias, dnode for dilation | 126 | #knode for kernel, bnode for bias, dnode for dilation, anode for activation |
| 77 | knode, bnode, dnode, activation = self.get_conv2d_params(scope_name) | 127 | knode, bnode, dnode, anode = self.get_conv2d_params(scope_name) |
| 78 | 128 | ||
| 79 | if dnode is not None: | 129 | if dnode is not None: |
| 80 | dilation = struct.unpack('i', dnode.attr['value'].tensor.tensor_content[0:4])[0] | 130 | dilation = struct.unpack('i', dnode.attr['value'].tensor.tensor_content[0:4])[0] |
| 81 | else: | 131 | else: |
| 82 | dilation = 1 | 132 | dilation = 1 |
| 83 | 133 | ||
| 134 | if anode is not None: | ||
| 135 | activation = anode.op | ||
| 136 | else: | ||
| 137 | activation = 'None' | ||
| 138 | |||
| 84 | padding = node.attr['padding'].s.decode("utf-8") | 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 | if dilation > 1 and scope_name + '/stack' in self.name_node_dict: | 141 | if dilation > 1 and scope_name + '/stack' in self.name_node_dict: |
| 87 | if self.name_node_dict[scope_name + '/stack'].op == "Const": | 142 | if self.name_node_dict[scope_name + '/stack'].op == "Const": |
| 88 | padding = 'SAME' | 143 | padding = 'SAME' |
| @@ -107,6 +162,15 @@ class TFConverter: · 18 unmodified lines | |||
| 107 | bias = btensor.tensor_content | 162 | bias = btensor.tensor_content |
| 108 | f.write(bias) | 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 | def dump_depth2space_to_file(self, node, f): | 175 | def dump_depth2space_to_file(self, node, f): |
| 112 | assert(node.op == 'DepthToSpace') | 176 | assert(node.op == 'DepthToSpace') |
| @@ -114,6 +178,9 @@ class TFConverter: · 1 unmodified line | |||
| 114 | block_size = node.attr['block_size'].i | 178 | block_size = node.attr['block_size'].i |
| 115 | np.array([self.op2code[node.op], block_size], dtype=np.uint32).tofile(f) | 179 | np.array([self.op2code[node.op], block_size], dtype=np.uint32).tofile(f) |
| 116 | self.converted_nodes.add(node.name) | 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 | def dump_mirrorpad_to_file(self, node, f): | 186 | def dump_mirrorpad_to_file(self, node, f): |
| @@ -127,6 +194,9 @@ class TFConverter: · 7 unmodified lines | |||
| 127 | paddings = pnode.attr['value'].tensor.tensor_content | 194 | paddings = pnode.attr['value'].tensor.tensor_content |
| 128 | f.write(paddings) | 195 | f.write(paddings) |
| 129 | self.converted_nodes.add(node.name) | 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 | def dump_layers_to_file(self, f): | 202 | def dump_layers_to_file(self, f): |
| @@ -147,10 +217,21 @@ class TFConverter: · 14 unmodified lines | |||
| 147 | self.dump_mirrorpad_to_file(node, f) | 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 | def dump_to_file(self): | 230 | def dump_to_file(self): |
| 151 | with open(self.outfile, 'wb') as f: | 231 | with open(self.outfile, 'wb') as f: |
| 152 | self.dump_layers_to_file(f) | 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 | def generate_name_node_dict(self): | 237 | def generate_name_node_dict(self): |
| @@ -212,19 +293,29 @@ class TFConverter: · 55 unmodified lines | |||
| 212 | return name[0:index] | 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 | for node in self.nodes: | 298 | for node in self.nodes: |
| 217 | if node.op == 'Conv2D': | 299 | if node.op == 'Conv2D': |
| 218 | scope = TFConverter.get_scope_name(node.name) | 300 | scope = TFConverter.get_scope_name(node.name) |
| 219 | self.conv2d_scope_names.add(scope) | 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 | def run(self): | 313 | def run(self): |
| 223 | self.generate_name_node_dict() | 314 | self.generate_name_node_dict() |
| 224 | self.generate_output_names() | 315 | self.generate_output_names() |
| 225 | self.remove_identity() | 316 | self.remove_identity() |
| 226 | self.generate_edges() | 317 | self.generate_edges() |
| 227 | self.generate_conv2d_scope_names() | 318 | self.generate_conv2d_scope_info() |
| 228 | 319 | ||
| 229 | if self.dump4tb: | 320 | if self.dump4tb: |
| 230 | self.dump_for_tensorboard() | 321 | self.dump_for_tensorboard() |