Saves a frozen tensorflow graph (a protobuf file). See also https://leimao.github.io/blog/Save-Load-Inference-From-TF-Frozen-Graph/ Parameters ---------- sess : tensorflow session session with graph to be saved checkpoint : string checkpoint of tensorflow
(sess, checkpoint, output, output_dir=None)
| 159 | |
| 160 | |
| 161 | def tf_to_pb(sess, checkpoint, output, output_dir=None): |
| 162 | """ |
| 163 | |
| 164 | Saves a frozen tensorflow graph (a protobuf file). |
| 165 | See also https://leimao.github.io/blog/Save-Load-Inference-From-TF-Frozen-Graph/ |
| 166 | |
| 167 | Parameters |
| 168 | ---------- |
| 169 | sess : tensorflow session |
| 170 | session with graph to be saved |
| 171 | |
| 172 | checkpoint : string |
| 173 | checkpoint of tensorflow model to be converted to protobuf (output will be <checkpoint>.pb) |
| 174 | |
| 175 | output : list of strings |
| 176 | list of the names of output nodes (is returned by load_models) |
| 177 | |
| 178 | output_dir : string, optional |
| 179 | path to the directory that exported models should be saved to. |
| 180 | If None, will export to the directory of the checkpoint file. |
| 181 | """ |
| 182 | |
| 183 | output_dir = os.path.expanduser(output_dir) if output_dir else os.path.dirname(checkpoint) |
| 184 | ckpt_base = os.path.basename(checkpoint) |
| 185 | |
| 186 | # save graph to pbtxt file |
| 187 | pbtxt_file = os.path.normpath(output_dir + "/" + ckpt_base + ".pbtxt") |
| 188 | tf.io.write_graph(sess.graph.as_graph_def(), "", pbtxt_file, as_text=True) |
| 189 | |
| 190 | # create frozen graph from pbtxt file |
| 191 | pb_file = os.path.normpath(output_dir + "/" + ckpt_base + ".pb") |
| 192 | frozen_graph_def = tf.compat.v1.graph_util.convert_variables_to_constants( |
| 193 | sess, |
| 194 | sess.graph_def, |
| 195 | output, |
| 196 | ) |
| 197 | with open(pb_file, "wb") as file: |
| 198 | file.write(frozen_graph_def.SerializeToString()) |
| 199 | |
| 200 | |
| 201 | def export_model( |