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 | 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() |