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.
Unified · Split
| @@ -33,9 +33,10 @@ class TFConverter: · 32 unmodified lines | ||
| 33 | 33 | self.output_names = [] |
| 34 | 34 | self.name_node_dict = {} |
| 35 | 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 | 37 | self.conv_paddings = {'VALID':0, 'SAME':1} |
| 38 | 38 | self.converted_nodes = set() |
| 39 | + self.conv2d_scope_names = set() | |
| 39 | 40 | self.op2code = {'Conv2D':1, 'DepthToSpace':2, 'MirrorPad':3} |
| 40 | 41 | self.mirrorpad_mode = {'CONSTANT':0, 'REFLECT':1, 'SYMMETRIC':2} |
| 41 | 42 | |
| @@ -47,30 +48,45 @@ class TFConverter: · 5 unmodified lines | ||
| 47 | 48 | print('graph saved, run "tensorboard --logdir=/tmp/graph" to see it') |
| 48 | 49 | |
| 49 | 50 | |
| 50 | − def get_conv2d_params(self, node): | |
| 51 | − knode = self.name_node_dict[node.input[1]] | |
| 52 | − bnode = None | |
| 53 | − activation = 'None' | |
| 54 | − next = self.edges[node.name][0] | |
| 55 | − if next.op == 'BiasAdd': | |
| 56 | − self.converted_nodes.add(next.name) | |
| 57 | − bnode = self.name_node_dict[next.input[1]] | |
| 58 | − next = self.edges[next.name][0] | |
| 59 | − if next.op in self.conv_activations: | |
| 60 | − self.converted_nodes.add(next.name) | |
| 61 | − activation = next.op | |
| 62 | − return knode, bnode, activation | |
| 51 | + def get_conv2d_params(self, conv2d_scope_name): | |
| 52 | + knode = self.name_node_dict[conv2d_scope_name + '/kernel'] | |
| 53 | + bnode = self.name_node_dict[conv2d_scope_name + '/bias'] | |
| 54 | + | |
| 55 | + if conv2d_scope_name + '/dilation_rate' in self.name_node_dict: | |
| 56 | + dnode = self.name_node_dict[conv2d_scope_name + '/dilation_rate'] | |
| 57 | + else: | |
| 58 | + dnode = None | |
| 59 | + | |
| 60 | + # 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 | |
| 62 | + if conv2d_scope_name + '/BiasAdd' in self.edges: | |
| 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 | 70 | def dump_conv2d_to_file(self, node, f): |
| 66 | 71 | assert(node.op == 'Conv2D') |
| 67 | 72 | self.layer_number = self.layer_number + 1 |
| 68 | 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] | |
| 72 | − padding = node.attr['padding'].s | |
| 73 | − padding = self.conv_paddings[padding.decode("utf-8")] | |
| 75 | + 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) | |
| 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 | 91 | ktensor = knode.attr['value'].tensor |
| 76 | 92 | filter_height = ktensor.tensor_shape.dim[0].size |
| @@ -126,9 +142,15 @@ class TFConverter: · 49 unmodified lines | ||
| 126 | 142 | for node in self.nodes: |
| 127 | 143 | if node.name in self.converted_nodes: |
| 128 | 144 | continue |
| 129 | − if node.op == 'Conv2D': | |
| 130 | − self.dump_conv2d_to_file(node, f) | |
| 131 | − elif node.op == 'DepthToSpace': | |
| 145 | + | |
| 146 | + # conv2d with dilation generates very complex nodes, so handle it in special | |
| 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 | 154 | self.dump_depth2space_to_file(node, f) |
| 133 | 155 | elif node.op == 'MirrorPad': |
| 134 | 156 | self.dump_mirrorpad_to_file(node, f) |
| @@ -192,11 +214,27 @@ class TFConverter: · 57 unmodified lines | ||
| 192 | 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 | 232 | def run(self): |
| 196 | 233 | self.generate_name_node_dict() |
| 197 | 234 | self.generate_output_names() |
| 198 | 235 | self.remove_identity() |
| 199 | 236 | self.generate_edges() |
| 237 | + self.generate_conv2d_scope_names() | |
| 200 | 238 | |
| 201 | 239 | if self.dump4tb: |
| 202 | 240 | self.dump_for_tensorboard() |