currently, the layer number is at the beginning of the .model file, so we have to scan twice in python script, the first scan to get the layer number. Only one scan needed after put the layer number at the end of .model file.
Showing the changes to tools/python/convert_from_tensorflow.py (1 of 2 files). Show all
Unified · Split
| @@ -129,15 +129,6 @@ class TFConverter: · 128 unmodified lines | |||
| 129 | self.converted_nodes.add(node.name) | 129 | self.converted_nodes.add(node.name) |
| 130 | 130 | ||
| 131 | 131 | ||
| 132 | def generate_layer_number(self): | ||
| 133 | # in current hard code implementation, the layer number is the first data written to the native model file | ||
| 134 | # it is not easy to know it at the beginning time in the general converter, so first do a dry run for compatibility | ||
| 135 | # will be refined later. | ||
| 136 | with open('/tmp/tmp.model', 'wb') as f: | ||
| 137 | self.dump_layers_to_file(f) | ||
| 138 | self.converted_nodes.clear() | ||
| 139 | |||
| 140 | |||
| 141 | def dump_layers_to_file(self, f): | 132 | def dump_layers_to_file(self, f): |
| 142 | for node in self.nodes: | 133 | for node in self.nodes: |
| 143 | if node.name in self.converted_nodes: | 134 | if node.name in self.converted_nodes: |
| @@ -157,10 +148,9 @@ class TFConverter: · 13 unmodified lines | |||
| 157 | 148 | ||
| 158 | 149 | ||
| 159 | def dump_to_file(self): | 150 | def dump_to_file(self): |
| 160 | self.generate_layer_number() | ||
| 161 | with open(self.outfile, 'wb') as f: | 151 | with open(self.outfile, 'wb') as f: |
| 162 | np.array([self.layer_number], dtype=np.uint32).tofile(f) | ||
| 163 | self.dump_layers_to_file(f) | 152 | self.dump_layers_to_file(f) |
| 153 | np.array([self.layer_number], dtype=np.uint32).tofile(f) | ||
| 164 | 154 | ||
| 165 | 155 | ||
| 166 | def generate_name_node_dict(self): | 156 | def generate_name_node_dict(self): |