conv2d with dilation > 1 generates tens of nodes in graph, it is not easy to parse each node one by one, so we do special tricks to parse the conv2d layer.
Showing the changes to tools/python (1 of 1 file). Show all
Unified · Split
| @@ -33,9 +33,10 @@ class TFConverter: · 32 unmodified lines | |||
| 33 | self.output_names = [] | 33 | self.output_names = [] |
| 34 | self.name_node_dict = {} | 34 | self.name_node_dict = {} |
| 35 | self.edges = {} | 35 | self.edges = {} |
| 36 | self.conv_activations = {'Relu':0, 'Tanh':1, 'Sigmoid':2, 'LeakyRelu':4} | 36 | self.conv_activations = {'Relu':0, 'Tanh':1, 'Sigmoid':2, 'None':3, 'LeakyRelu':4} |
| 37 | self.conv_paddings = {'VALID':0, 'SAME':1} | 37 | self.conv_paddings = {'VALID':0, 'SAME':1} |
| 38 | self.converted_nodes = set() | 38 | self.converted_nodes = set() |
| 39 | self.conv2d_scope_names = set() | ||
| 39 | self.op2code = {'Conv2D':1, 'DepthToSpace':2, 'MirrorPad':3} | 40 | self.op2code = {'Conv2D':1, 'DepthToSpace':2, 'MirrorPad':3} |
| 40 | self.mirrorpad_mode = {'CONSTANT':0, 'REFLECT':1, 'SYMMETRIC':2} | 41 | self.mirrorpad_mode = {'CONSTANT':0, 'REFLECT':1, 'SYMMETRIC':2} |
| 41 | 42 | ||
| @@ -47,30 +48,45 @@ class TFConverter: · 5 unmodified lines | |||
| 47 | print('graph saved, run "tensorboard --logdir=/tmp/graph" to see it') | 48 | print('graph saved, run "tensorboard --logdir=/tmp/graph" to see it') |
| 48 | 49 | ||
| 49 | 50 | ||
| 50 | def get_conv2d_params(self, node): | 51 | def get_conv2d_params(self, conv2d_scope_name): |
| 51 | knode = self.name_node_dict[node.input[1]] | 52 | knode = self.name_node_dict[conv2d_scope_name + '/kernel'] |
| 52 | bnode = None | 53 | bnode = self.name_node_dict[conv2d_scope_name + '/bias'] |
| 53 | activation = 'None' | 54 | |
| 54 | next = self.edges[node.name][0] | 55 | if conv2d_scope_name + '/dilation_rate' in self.name_node_dict: |
| 55 | if next.op == 'BiasAdd': | 56 | dnode = self.name_node_dict[conv2d_scope_name + '/dilation_rate'] |
| 56 | self.converted_nodes.add(next.name) | 57 | else: |
| 57 | bnode = self.name_node_dict[next.input[1]] | 58 | dnode = None |
| 58 | next = self.edges[next.name][0] | 59 | |
| 59 | if next.op in self.conv_activations: | 60 | # the BiasAdd name is possible be changed into the output name, |
| 60 | self.converted_nodes.add(next.name) | 61 | # if activation is None, and BiasAdd.next is the last op which is Identity |
| 61 | activation = next.op | 62 | if conv2d_scope_name + '/BiasAdd' in self.edges: |
| 62 | return knode, bnode, activation | 63 | activation = self.edges[conv2d_scope_name + '/BiasAdd'][0] |
| 64 | activation = activation.op | ||
| 65 | else: | ||
| 66 | activation = 'None' | ||
| 67 | return knode, bnode, dnode, activation | ||
| 63 | 68 | ||
| 64 | 69 | ||
| 65 | def dump_conv2d_to_file(self, node, f): | 70 | def dump_conv2d_to_file(self, node, f): |
| 66 | assert(node.op == 'Conv2D') | 71 | assert(node.op == 'Conv2D') |
| 67 | self.layer_number = self.layer_number + 1 | 72 | self.layer_number = self.layer_number + 1 |
| 68 | self.converted_nodes.add(node.name) | 73 | self.converted_nodes.add(node.name) |
| 69 | knode, bnode, activation = self.get_conv2d_params(node) | ||
| 70 | 74 | ||
| 71 | dilation = node.attr['dilations'].list.i[0] | 75 | scope_name = TFConverter.get_scope_name(node.name) |
| 72 | padding = node.attr['padding'].s | 76 | #knode for kernel, bnode for bias, dnode for dilation |
| 73 | padding = self.conv_paddings[padding.decode("utf-8")] | 77 | knode, bnode, dnode, activation = self.get_conv2d_params(scope_name) |
| 78 | |||
| 79 | if dnode is not None: | ||
| 80 | dilation = struct.unpack('i', dnode.attr['value'].tensor.tensor_content[0:4])[0] | ||
| 81 | else: | ||
| 82 | dilation = 1 | ||
| 83 | |||
| 84 | 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. | ||
| 86 | if dilation > 1 and scope_name + '/stack' in self.name_node_dict: | ||
| 87 | if self.name_node_dict[scope_name + '/stack'].op == "Const": | ||
| 88 | padding = 'SAME' | ||
| 89 | padding = self.conv_paddings[padding] | ||
| 74 | 90 | ||
| 75 | ktensor = knode.attr['value'].tensor | 91 | ktensor = knode.attr['value'].tensor |
| 76 | filter_height = ktensor.tensor_shape.dim[0].size | 92 | filter_height = ktensor.tensor_shape.dim[0].size |
| @@ -126,9 +142,15 @@ class TFConverter: · 49 unmodified lines | |||
| 126 | for node in self.nodes: | 142 | for node in self.nodes: |
| 127 | if node.name in self.converted_nodes: | 143 | if node.name in self.converted_nodes: |
| 128 | continue | 144 | continue |
| 129 | if node.op == 'Conv2D': | 145 | |
| 130 | self.dump_conv2d_to_file(node, f) | 146 | # conv2d with dilation generates very complex nodes, so handle it in special |
| 131 | elif node.op == 'DepthToSpace': | 147 | scope_name = TFConverter.get_scope_name(node.name) |
| 148 | if scope_name in self.conv2d_scope_names: | ||
| 149 | if node.op == 'Conv2D': | ||
| 150 | self.dump_conv2d_to_file(node, f) | ||
| 151 | continue | ||
| 152 | |||
| 153 | if node.op == 'DepthToSpace': | ||
| 132 | self.dump_depth2space_to_file(node, f) | 154 | self.dump_depth2space_to_file(node, f) |
| 133 | elif node.op == 'MirrorPad': | 155 | elif node.op == 'MirrorPad': |
| 134 | self.dump_mirrorpad_to_file(node, f) | 156 | self.dump_mirrorpad_to_file(node, f) |
| @@ -192,11 +214,27 @@ class TFConverter: · 57 unmodified lines | |||
| 192 | self.edges[input] = [node] | 214 | self.edges[input] = [node] |
| 193 | 215 | ||
| 194 | 216 | ||
| 217 | @staticmethod | ||
| 218 | def get_scope_name(name): | ||
| 219 | index = name.rfind('/') | ||
| 220 | if index == -1: | ||
| 221 | return "" | ||
| 222 | return name[0:index] | ||
| 223 | |||
| 224 | |||
| 225 | def generate_conv2d_scope_names(self): | ||
| 226 | for node in self.nodes: | ||
| 227 | if node.op == 'Conv2D': | ||
| 228 | scope = TFConverter.get_scope_name(node.name) | ||
| 229 | self.conv2d_scope_names.add(scope) | ||
| 230 | |||
| 231 | |||
| 195 | def run(self): | 232 | def run(self): |
| 196 | self.generate_name_node_dict() | 233 | self.generate_name_node_dict() |
| 197 | self.generate_output_names() | 234 | self.generate_output_names() |
| 198 | self.remove_identity() | 235 | self.remove_identity() |
| 199 | self.generate_edges() | 236 | self.generate_edges() |
| 237 | self.generate_conv2d_scope_names() | ||
| 200 | 238 | ||
| 201 | if self.dump4tb: | 239 | if self.dump4tb: |
| 202 | self.dump_for_tensorboard() | 240 | self.dump_for_tensorboard() |