首页 > 编程 > Python > 正文

TensorFlow:将ckpt文件固化成pb文件教程

2020-02-15 21:15:40
字体:
来源:转载
供稿:网友

本文是将yolo3目标检测框架训练出来的ckpt文件固化成pb文件,主要利用了GitHub上的该项目。

为什么要最终生成pb文件呢?简单来说就是直接通过tf.saver保存行程的ckpt文件其变量数据和图是分开的。我们知道TensorFlow是先画图,然后通过placeholde往图里面喂数据。这种解耦形式存在的方法对以后的迁移学习以及对程序进行微小的改动提供了极大的便利性。但是对于训练好,以后不再改变的话这种存在就不再需要。一方面,ckpt文件储存的数据都是变量,既然我们不再改动,就应当让其变成常量,直接‘烧'到图里面。另一方面,对于线上的模型,我们一般是通过C++或者C语言编写的程序进行调用。所以一般模型最终形式都是应该写成pb文件的形式。

由于这次的程序直接从GitHub上下载后改动较小就能够运行,也就是自己写了很少一部分程序。因此进行调试的时候还出现了以前根本没有注意的一些小问题,同时发现自己对TensorFlow还需要更加详细的去研读。

首先对程序进行保存的时候,利用 saver = tf.train.Saver(), saver.save(sess,checkpoint_path,global_step=global_step)对训练的数据进行保存,保存格式为ckpt。但是在恢复的时候一直提示有问题,(其恢复语句为:saver = tf.train.Saver(), saver.restore(sess,ckpt_path),其中,ckpt_path是保存ckpt的文件夹路径)。出现问题的原因我估计是因为我是按照每50个epoch进行保存,而不是让其进行固定次数的batch进行保存,这种固定batch次数的保存系统会自动保存最近5次的ckpt文件(该方法的ckpt_path=tf.train,latest_checkpoint('ckpt/')进行回复)。那么如何将利用epoch的次数进行保存呢(这种保存不是近5次的保存,而是每进行一次保存就会留下当时保存的ckpt,而那种按照batch的会在第n次保存,会将n-5次的删除,n>5)。

我们可以利用:ckpt = tf.train.get_checkpoint_state(ckpt_path),获取最新的ckptpoint文件,然后利用saver.restore(sess,ckpt.checkpoint_path)进行恢复。当然为了安全起见,应该对ckpt和ckpt.checkpoint_path进行判断是否存在后,再进行恢复语句的调用,建议打开ckptpoint看一下,里面记录的最近五次的model的路径,一目了然。即:

  saver = tf.train.Saver()  ckpt = tf.train.get_checkpoint_state(model_path)  if ckpt and ckpt.model_checkpoint_path:    saver.restore(sess, ckpt.model_checkpoint_path)

对于固化网络,网上有很多的介绍。之所以再介绍,还是由于是用了别人的网络而不是自己的网络遇到的坑。在固化时候我们需要知道输出tensor的名字,而再恢复的时候我们需要知道placeholder的名字。但是,如果网络复杂或者别人的网络命名比较复杂,或者name=,根本就没有自己命名而用的系统自定义的,这样捋起来还是比较费劲的。当时在网上查找的一些方法,像打印整个网络变量的方法(先不管输出的网路名称,甚至随便起一个名字,先固化好pb文件,然后对pb文件进行读取,最后打印操作的名字:

发表评论 共有条评论
用户名: 密码:
验证码: 匿名发表